From 14d70d98703f71bcbbc125b3e90a3ad5060c703a Mon Sep 17 00:00:00 2001 From: bjf-frz Date: Thu, 18 Jun 2026 16:54:05 +0800 Subject: [PATCH 01/40] [Perf]Enable regional compile for Cosmos3 HSDP blocks (#4485) Signed-off-by: bjf-frz --- tests/diffusion/test_compile.py | 35 +++++++++++++++++++++++++++++++++ vllm_omni/diffusion/compile.py | 33 +++++++++++++++++++++++++++++-- 2 files changed, 66 insertions(+), 2 deletions(-) create mode 100644 tests/diffusion/test_compile.py diff --git a/tests/diffusion/test_compile.py b/tests/diffusion/test_compile.py new file mode 100644 index 000000000..522ef343f --- /dev/null +++ b/tests/diffusion/test_compile.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch.nn as nn + +from vllm_omni.diffusion.compile import regionally_compile + + +class _WrappedBlock(nn.Module): + def __init__(self) -> None: + super().__init__() + self.compile_called = False + + def compile(self, *args, **kwargs): + self.compile_called = True + return self + + +class _ModelWithWrappedRepeatedBlocks(nn.Module): + _repeated_blocks = ["OriginalBlock"] + _layerwise_offload_blocks_attrs = ["transformer_blocks"] + + def __init__(self) -> None: + super().__init__() + self.transformer_blocks = nn.ModuleList([_WrappedBlock(), _WrappedBlock()]) + self.other_blocks = nn.ModuleList([_WrappedBlock()]) + + +def test_regionally_compile_matches_wrapped_blocks_by_declared_container_attr(): + model = _ModelWithWrappedRepeatedBlocks() + + regionally_compile(model) + + assert all(block.compile_called for block in model.transformer_blocks) + assert not model.other_blocks[0].compile_called diff --git a/vllm_omni/diffusion/compile.py b/vllm_omni/diffusion/compile.py index 66a4675f3..0e4a3ff50 100644 --- a/vllm_omni/diffusion/compile.py +++ b/vllm_omni/diffusion/compile.py @@ -8,6 +8,25 @@ from vllm.logger import init_logger logger = init_logger(__name__) +def _matches_repeated_block( + name: str, + module: nn.Module, + repeated_blocks: list[str], + repeated_block_attrs: list[str], +) -> bool: + class_name = module.__class__.__name__ + if class_name in repeated_blocks: + return True + + for attr in ("_fsdp_wrapped_module", "module", "_orig_mod"): + wrapped = getattr(module, attr, None) + if wrapped is not None and wrapped.__class__.__name__ in repeated_blocks: + return True + + parts = name.split(".") + return len(parts) >= 2 and parts[-2] in repeated_block_attrs and parts[-1].isdigit() + + def regionally_compile(model: nn.Module, *compile_args: Any, **compile_kwargs: Any) -> nn.Module: """ Apply regional compilation to a PyTorch model. @@ -27,15 +46,25 @@ def regionally_compile(model: nn.Module, *compile_args: Any, **compile_kwargs: A logger.warning("Regional compilation skipped because the model does not define `_repeated_blocks`.") return model + repeated_block_attrs = getattr(model, "_layerwise_offload_blocks_attrs", []) + # Check if we have modules with the specified class names has_compiled_region = False - for submod in model.modules(): - if submod.__class__.__name__ in repeated_blocks: + compiled_region_count = 0 + for name, submod in model.named_modules(): + if _matches_repeated_block(name, submod, repeated_blocks, repeated_block_attrs): # Compile this submodule submod.compile(*compile_args, **compile_kwargs) has_compiled_region = True + compiled_region_count += 1 if not has_compiled_region: logger.warning(f"Regional compilation skipped because {repeated_blocks} classes are not found in the model.") + else: + logger.info( + "Regional compilation applied to %d module(s) for repeated blocks %s.", + compiled_region_count, + repeated_blocks, + ) return model From db8bc2c71d76ddbe475eb291b8bcaab09ca20d02 Mon Sep 17 00:00:00 2001 From: Alicia <115451386+congw729@users.noreply.github.com> Date: Thu, 18 Jun 2026 16:56:13 +0800 Subject: [PATCH 02/40] Update VoxCPM and Qwen3TTS perf baseline. (#4521) Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com> --- tests/dfx/perf/tests/test_tts.json | 42 +++++++++++++------------- tests/dfx/perf/tests/test_voxcpm2.json | 26 ++++++++-------- 2 files changed, 34 insertions(+), 34 deletions(-) diff --git a/tests/dfx/perf/tests/test_tts.json b/tests/dfx/perf/tests/test_tts.json index ba1f95261..1ba273a2c 100644 --- a/tests/dfx/perf/tests/test_tts.json +++ b/tests/dfx/perf/tests/test_tts.json @@ -27,10 +27,10 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 122.4504 + 139.9434 ], "median_audio_rtf": [ - 0.1472 + 0.1683 ] } }, @@ -55,17 +55,17 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 199.1197, + 227.5653, 647.1272, 4777.2714 ], "median_audio_rtf": [ - 0.2899, + 0.3313, 0.3642, 1.2282 ], "audio_throughput": [ - 20.5634, + 23.501, 40.634, 43.0758 ] @@ -129,10 +129,10 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 43.5669 + 49.7908 ], "median_audio_rtf": [ - 0.1247 + 0.1425 ] } }, @@ -162,17 +162,17 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 65.9174, + 75.3341, 713.7727, 5941.7179 ], "median_audio_rtf": [ - 0.1626, + 0.1859, 0.3148, 1.1662 ], "audio_throughput": [ - 31.9068, + 36.465, 49.6268, 48.3091 ] @@ -228,13 +228,13 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 49.6105 + 56.6977 ], "median_audio_rtf": [ - 0.1361 + 0.1555 ], "audio_throughput": [ - 10.0881 + 11.5292 ] } }, @@ -259,10 +259,10 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 45.5391 + 52.0447 ], "median_audio_rtf": [ - 0.1271 + 0.1452 ] } }, @@ -291,17 +291,17 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 66.1471, + 75.5967, 587.5877, 4960.3726 ], "median_audio_rtf": [ - 0.1637, + 0.1871, 0.3143, 1.1703 ], "audio_throughput": [ - 35.1923, + 40.2198, 48.9046, 48.0505 ] @@ -328,13 +328,13 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 49.3466 + 56.3961 ], "median_audio_rtf": [ - 0.1349 + 0.1542 ], "audio_throughput": [ - 9 + 9.6672 ] } } diff --git a/tests/dfx/perf/tests/test_voxcpm2.json b/tests/dfx/perf/tests/test_voxcpm2.json index 0aaca446f..beb40d5d3 100644 --- a/tests/dfx/perf/tests/test_voxcpm2.json +++ b/tests/dfx/perf/tests/test_voxcpm2.json @@ -23,10 +23,10 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 113.9351 + 130.2115 ], "median_audio_rtf": [ - 0.0427 + 0.0488 ] } }, @@ -47,13 +47,13 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 415.9187 + 475.3357 ], "median_audio_rtf": [ - 0.1462 + 0.1671 ], "audio_throughput": [ - 20.0128 + 22.8717 ] } }, @@ -100,10 +100,10 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 66.2748 + 75.7427 ], "median_audio_rtf": [ - 0.0366 + 0.0419 ] } }, @@ -127,13 +127,13 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 307.6373 + 351.5855 ], "median_audio_rtf": [ - 0.1056 + 0.1207 ], "audio_throughput": [ - 64.5428 + 73.7632 ] } }, @@ -157,13 +157,13 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 70.0958 + 80.1095 ], "median_audio_rtf": [ - 0.0424 + 0.0485 ], "audio_throughput": [ - 16.075 + 18.3715 ] } }, From c78274613369a2244c002e8e17c9065e0e60e94e Mon Sep 17 00:00:00 2001 From: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com> Date: Thu, 18 Jun 2026 16:58:40 +0800 Subject: [PATCH 03/40] Add NPU support for OmniMRotaryEmbedding (#3609) Signed-off-by: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com> --- .../model_executor/layers/rotary_embedding/mrope.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/vllm_omni/model_executor/layers/rotary_embedding/mrope.py b/vllm_omni/model_executor/layers/rotary_embedding/mrope.py index 3d3a88d87..9c65a3c99 100644 --- a/vllm_omni/model_executor/layers/rotary_embedding/mrope.py +++ b/vllm_omni/model_executor/layers/rotary_embedding/mrope.py @@ -16,12 +16,18 @@ import itertools import torch from transformers import PretrainedConfig from vllm.logger import init_logger -from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding + +from vllm_omni.platforms import current_omni_platform + +if current_omni_platform.is_npu(): + from vllm_ascend.ops.rotary_embedding import AscendMRotaryEmbedding as _BaseMRotaryEmbedding +else: + from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding as _BaseMRotaryEmbedding logger = init_logger(__name__) -class OmniMRotaryEmbedding(MRotaryEmbedding): +class OmniMRotaryEmbedding(_BaseMRotaryEmbedding): """Omni-extended MRotaryEmbedding with multimodal position computation. Extends the upstream MRotaryEmbedding with additional class methods for From dff8cda84425cb50cdb2504dc253a960ccbc6317 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Thu, 18 Jun 2026 17:01:24 +0800 Subject: [PATCH 04/40] [SKILLS] Add vllm-omni-test agent skill for CI-aligned test generation (#4434) Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: [Your Name] <[Your Email]> --- .claude/skills/add-diffusion-model/SKILL.md | 32 +- .claude/skills/add-tts-model/SKILL.md | 74 +- .claude/skills/readme.md | 6 + .claude/skills/vllm-omni-test/SKILL.md | 890 ++++++++++++++++++ .../vllm-omni-test/references/test-routing.md | 372 ++++++++ docs/contributing/ci/CI_5levels.md | 62 +- 6 files changed, 1419 insertions(+), 17 deletions(-) create mode 100644 .claude/skills/vllm-omni-test/SKILL.md create mode 100644 .claude/skills/vllm-omni-test/references/test-routing.md diff --git a/.claude/skills/add-diffusion-model/SKILL.md b/.claude/skills/add-diffusion-model/SKILL.md index c080ca412..c7cdb47ac 100644 --- a/.claude/skills/add-diffusion-model/SKILL.md +++ b/.claude/skills/add-diffusion-model/SKILL.md @@ -288,9 +288,36 @@ Required updates: 5. `examples/offline_inference/xxx/README.md` — offline example docs 6. `examples/online_serve/xxx/README.md` — online serve docs -### Step 8: Add E2E Tests (Recommended) +### Step 8: Add E2E Tests -Create `tests/e2e/online_serving/test_your_model_expansion.py`. +**Follow the [vllm-omni-test skill](../vllm-omni-test/SKILL.md)** for markers, file naming, Buildkite wiring, and run commands. Also read [l4_functionality_tests.inc.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md) and [CI_5levels.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/CI_5levels.md). + +Classify the model's **CI priority** first: + +| Priority | Required test levels | Files & markers | +|----------|---------------------|-----------------| +| **High** (listed in [#1832](https://github.com/vllm-project/vllm-omni/issues/1832) or on the diffusion hot path) | **L1** · **L2** online · **L3** online + offline · **L4** feature + performance | See table below | +| **Medium** (*normal priority* in L4 docs) | **L3** online + offline · **L4** feature only | Fewer L4 parametrized rows | +| **Low** | **L4** feature only | One or two `*_expansion.py` cases | + +**Per-level deliverables (diffusion / `pytest.mark.diffusion`):** + +| Level | Location | Marker | CI pipeline | Notes | +|-------|----------|--------|-------------|-------| +| **L1** | `tests/diffusion/models/{slug}/`, `tests/diffusion/cache/`, transformer unit tests | `core_model` + `cpu` | `test-ready.yml` | Weight remap, `_sp_plan`, cache enabler registration, shape contracts | +| **L2** | `tests/e2e/online_serving/test_{slug}.py` (and offline if the category is offline-first) | **`core_model` + `advanced_model`** (both on baseline smoke) + `diffusion` + `@hardware_test` / `hardware_marks` | `test-ready.yml` | Default deploy smoke — minimal `num_inference_steps`, single prompt | +| **L3** | `tests/e2e/online_serving/test_{slug}.py` **and** `tests/e2e/offline_inference/test_{slug}.py` when offline matters | Baseline smoke: **`core_model` + `advanced_model`**; heavier cases: `advanced_model` only (+ `diffusion`) | `test-merge.yml` **or** merged into nightly diffusion function job | Real weights, streaming/API paths, LoRA/offload smoke | +| **L4** | `tests/e2e/online_serving/test_{slug}_expansion.py` (+ offline expansion if needed) | `full_model` + `diffusion` | `test-nightly.yml` (X2I / X2V / X2A function groups) | Feature combos per [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); perf → `tests/dfx/perf/tests/test_{model}_vllm_omni.json` | + +**L2 & L3 online — same file, dual marks on the baseline smoke:** The **first / simplest** case in `test_{slug}.py` (default deploy, minimal steps, single prompt) should carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** function so L2 (`test-ready.yml`) and L3 (`test-merge.yml`) share one smoke test. Heavier deploy variants or API paths in the same file use **`advanced_model` only**. When L3 moves to nightly, migrate those heavier cases into `test_{slug}_expansion.py` with `full_model` and remove the dedicated `test-merge.yml` job (see `test_longcat_image_expansion.py`, `test_qwen_image_expansion.py`). + +**L4 design (high priority):** Combine multiple supported features (Cache-DiT, TP, USP, CFG, HSDP, CPU offload, quantization) into **few parametrized** `OmniServerParams` rows so each feature appears in at least one case without exploding GPU jobs. Shard single-GPU vs multi-GPU cases across the nightly X2I/X2V function steps (1-GPU vs `distributed_cuda`). + +**L4 design (medium / low):** One or two parametrized rows covering the best quality/perf trade-off; skip perf JSON unless the model is high priority. + +**Reference implementations:** `tests/e2e/online_serving/test_qwen_image_edit_expansion.py`, `tests/e2e/online_serving/test_longcat_image_expansion.py`, `tests/e2e/online_serving/test_hunyuan_video_15_expansion.py`. + +**Keep model-specific code inside test modules — not `tests/helpers/{slug}.py`:** deploy constants, prompts, sampling dicts, and inline `request_config` / `form_data` belong in each `test_{slug}.py` and `test_{slug}_expansion.py`. Do not add per-model files under `tests/helpers/`; reuse only repo-wide harness (`mark`, `media`, `runtime`, `stage_config`, `assertions`). **L2+ online/offline e2e:** reuse or add `send_*_request` in `tests/helpers/runtime.py` — tests call the handler, not raw `omni.generate` / HTTP. See [vllm-omni-test skill](../vllm-omni-test/SKILL.md) § **Runtime send helpers**. ### Step 9: Add Cache-DiT Acceleration @@ -568,6 +595,7 @@ See the [Profiling Single-Stage Diffusion](../../../docs/contributing/profiling. ## Reference Files +- [vllm-omni-test skill](../vllm-omni-test/SKILL.md) — L1–L4 markers, naming, Buildkite wiring, run commands - [Transformer Adaptation](references/transformer-adaptation.md) — porting transformers from diffusers - [Custom Model Patterns](references/custom-model-patterns.md) — patterns for non-diffusers models - [Parallelism Patterns](references/parallelism-patterns.md) — TP, SP/USP, CFG parallel, HSDP implementation details diff --git a/.claude/skills/add-tts-model/SKILL.md b/.claude/skills/add-tts-model/SKILL.md index 811a359d0..260a7e924 100644 --- a/.claude/skills/add-tts-model/SKILL.md +++ b/.claude/skills/add-tts-model/SKILL.md @@ -305,13 +305,72 @@ def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list: ] ``` +### Test Case Writing (CI Levels) + +**Follow the [vllm-omni-test skill](../vllm-omni-test/SKILL.md)** for markers, file naming (`test_{slug}.py` / `test_{slug}_expansion.py`), Buildkite wiring, and copy-paste run commands. Also read [CI_5levels.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/CI_5levels.md) and [tests_style.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/tests_style.md). + +Classify the model's **CI priority** first (high / medium / low). High-priority TTS models are typically those on the integration hot path or listed in tracking issues such as [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); medium and low tiers cover the long tail. When unsure, ask the reviewer which tier applies. + +| Priority | Required test levels | Files & markers | +|----------|---------------------|-----------------| +| **High** | **L1** unit/logic · **L2** online smoke · **L3** online + offline integration · **L4** feature + performance | See table below | +| **Medium** | **L3** online + offline · **L4** feature only | Skip dedicated L1/L2 unless fixing a logic bug | +| **Low** | **L4** feature only | One or two `*_expansion.py` parametrized cases | + +**Per-level deliverables (TTS / `pytest.mark.tts`):** + +| Level | Location | Marker | CI pipeline | Notes | +|-------|----------|--------|-------------|-------| +| **L1** | `tests/model_executor/…`, `tests/entrypoints/openai_api/…`, stage-processor tests | `core_model` + `cpu` | `test-ready.yml` | Prompt assembly, async_chunk helpers, `serving_speech` validation — no GPU | +| **L2** | `tests/e2e/online_serving/test_{slug}.py` | **`core_model` + `advanced_model`** (both on baseline smoke) + `tts` + `@hardware_test(...)` | `test-ready.yml` (`ready` label) | Default deploy smoke: single `/v1/audio/speech` or offline `OmniRunner` path | +| **L3** | `tests/e2e/online_serving/test_{slug}.py` **and** `tests/e2e/offline_inference/test_{slug}.py` | Baseline smoke: **`core_model` + `advanced_model`**; heavier cases: `advanced_model` only (+ `tts`) | `test-merge.yml` **or** merged into nightly TTS function job | Streaming, voice clone, batch/queue, async_chunk | +| **L4** | `tests/e2e/online_serving/test_{slug}_expansion.py`, optional offline expansion | `full_model` + `tts` | `test-nightly.yml` (`:full_moon: TTS · Function Test with L4`) | Feature matrix; perf → `tests/dfx/perf/tests/test_tts.json` | + +**L2 & L3 online — same file, dual marks on the baseline smoke:** The **first / simplest** case in `test_{slug}.py` (default deploy, single non-streaming `/v1/audio/speech` or equivalent offline path) should carry **both** `@pytest.mark.core_model` **and** `@pytest.mark.advanced_model` on the **same** function so it runs in L2 (`test-ready.yml`, `--run-level core_model`, basic validation) and L3 (`test-merge.yml`, `--run-level advanced_model`, deeper validation) without duplicating the test. In-tree examples: `test_voxcpm2_tts.py::test_text_to_audio_001`, `test_qwen3_tts_customvoice.py::test_text_to_audio_001`. + +Heavier scenarios in the same file use **`advanced_model` only** (streaming, extra languages, concurrency, async_chunk, batch). Example: `test_voice_clone_en_streaming_001` → `advanced_model` only. When migrating L3 to nightly, move those heavier cases into `test_{slug}_expansion.py` with `full_model` and drop the dedicated merge job (see `test_ming_tts_expansion.py`, `test_glm_tts_expansion.py`). + +```python +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_voice_clone_en_non_streaming_001(omni_server, openai_client) -> None: + openai_client.send_audio_speech_request({...}) +``` + +**L4 consolidation:** Prefer parametrized `OmniServerParams` rows (`default`, `async_chunk`, feature flags) in one expansion module rather than many merge-only files ([#1832](https://github.com/vllm-project/vllm-omni/issues/1832)). + +**L4 performance (high-priority models):** Add latency / throughput / stress rows in `tests/dfx/perf/tests/test_tts.json`, or a dedicated `tests/dfx/perf/tests/test_{slug}.json` when the model must not join the shared nightly server matrix before integration lands (see VoxCPM2 / Coqui XTTS pattern). Register the model in `benchmarks/tts/model_configs.yaml` for local `bench_tts.py`. Wire a **separate** `test-nightly.yml` Perf Test step when the JSON is not merged into `test_tts.json` yet. + +**Keep model-specific code inside test modules — not `tests/helpers/{slug}.py`:** + +- Put `MODEL`, deploy path, vendored `REF_AUDIO_URL`, `get_prompt()`, and inline `request_config` dicts **in each** `test_{slug}.py`, `test_{slug}_expansion.py`, offline `test_{slug}.py`, and L1 `test_{slug}_*.py` as needed. +- **Do not** add `tests/helpers/{slug}.py` (or `tests/helpers/{model_name}.py`) to deduplicate constants or request builders across those files. A little duplication is intentional; follow in-tree references such as `tests/e2e/online_serving/test_glm_tts.py` and `tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py`. +- `tests/helpers/` is for **repo-wide** harness code only (`mark.py`, `media.py`, `runtime.py`, `stage_config.py`, `assertions.py`, `fixtures/`). Import those; do not extend the tree with per-model modules. + +**Runtime send helpers (`tests/helpers/runtime.py`) — online and offline e2e:** + +| Path | Fixture | Call | +|------|---------|------| +| Online `/v1/*` | `openai_client` | `openai_client.send_*_request(request_config)` | +| Offline inference | `omni_runner_handler` | `omni_runner_handler.send_*_request(request_config)` | + +1. **Grep `runtime.py` first** — reuse `send_omni_request`, `send_diffusion_request`, `send_audio_speech_request` (online + offline Qwen-style TTS), `send_single_stage_tts_request` (Coqui XTTS / MOSS-TTS-Nano offline), etc. +2. **No matching helper** → add `send__request` (or `send__http_request` for negative/dfx) in **`runtime.py`** with general `assert_*` bundled inside, **then** call it from the test. +3. **Test file** holds `request_config` dicts only — not `omni.generate`, not `_collect_audio()`, not raw HTTP/SDK. +4. Wire `tests/helpers/runtime.py` into Buildkite `source_file_dependencies` when you add helpers. + +See [vllm-omni-test skill](../vllm-omni-test/SKILL.md) § **Runtime send helpers** for full tables and exceptions. + ### Deliverables - Updated `serving_speech.py` with all 5 integration points (single commit) - Client scripts and server launcher under `examples/online_serving/text_to_speech//` - Gradio demo with streaming and voice cloning UI in the same dir -- E2E online serving test (`tests/e2e/online_serving/test_.py`) -- Buildkite CI entry in `.buildkite/test-merge.yml` +- E2E tests per **Test Case Writing (CI Levels)** above (priority tier determines L1–L4 scope) +- Buildkite wired per level: `test-ready.yml` (L1/L2), `test-merge.yml` or nightly function job (L3), `test-nightly.yml` (L4) — see [vllm-omni-test skill](../vllm-omni-test/SKILL.md) - New section in `examples/online_serving/text_to_speech/README.md` (table row + per-model section). Do **not** create a top-level `examples/online_serving//` dir or a per-model `README.md` inside `text_to_speech//`. ### E2E test pitfalls to avoid @@ -326,9 +385,9 @@ def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list: - **Use the harness readiness gate.** The fixture waits for HTTP 200 on `/health`; don't add `time.sleep` in tests. If warmup is incomplete, make `/health` return non-200 until you're actually ready. -- **Mark with `@pytest.mark.core_model` + `hardware_test(res={"cuda": "H100"})`** - so the test lands in `test-ready.yml` (triggered by the `ready` label) rather - than only nightly. +- **Mark tests per the CI Levels table** — baseline smoke: `core_model` + `advanced_model`; heavier cases: `advanced_model` only; L4 expansion: `full_model` +- **No per-model helper modules** — do not create `tests/helpers/{slug}.py`; keep constants and `request_config` payloads in the test file +- **Online and offline e2e go through `runtime.py`** — `openai_client.send_audio_speech_request` (online); `omni_runner_handler.send_audio_speech_request` (Qwen-style offline) or `send_single_stage_tts_request` (single-stage offline). Add a new `send_*_request` in `runtime.py` when none fits; do not embed `omni.generate` or HTTP in tests ## Phase 4: Async Chunk (Streaming) @@ -461,8 +520,8 @@ Use this checklist when integrating a new TTS model: - [ ] Voice cloning works (if supported) - [ ] All response formats work (wav, mp3, flac, pcm) - [ ] Client scripts and server launcher created -- [ ] E2E online serving test written (`tests/e2e/online_serving/test_.py`) -- [ ] Buildkite CI entry added to `.buildkite/test-merge.yml` +- [ ] E2E tests added per model priority tier (see **Test Case Writing (CI Levels)**) +- [ ] Buildkite entries match level: `test-ready.yml` / `test-merge.yml` or nightly TTS job / `test-nightly.yml` - [ ] Gradio demo working - [ ] Documentation added (offline + online docs, nav, supported models) @@ -500,5 +559,6 @@ Project docs and adjacent skills: - [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) — supported models and usage - [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) — complete example of Phases 1–3 - [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) — complete example of all 5 phases +- [vllm-omni-test skill](../vllm-omni-test/SKILL.md) — L1–L4 markers, naming, Buildkite wiring, run commands - [Adding a TTS model (developer guide)](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/model/adding_tts_model.md) - `plan/voxcpm2_native_ar_design.md` — VoxCPM2's vLLM-native AR + side-computation pattern (distinct from the generator-based single-stage described above) diff --git a/.claude/skills/readme.md b/.claude/skills/readme.md index f42fcb6ab..b8e9bbb13 100644 --- a/.claude/skills/readme.md +++ b/.claude/skills/readme.md @@ -30,6 +30,12 @@ include: - `precheck-pr`: self-check a branch before creating a PR — validates PR title format, catches dead code, verifies accuracy/perf claims, and confirms merge readiness +- `vllm-omni-test`: guides generation and execution of CI-aligned tests (L1–L4), + pytest marker selection (`core_model` / `advanced_model` / `full_model`, + `omni` / `tts` / `diffusion`), Buildkite wiring (`test-ready.yml`, + `test-merge.yml`, `test-nightly.yml`, `test-weekly.yml`), and copy-paste + local plus CI-like `pytest` commands; see `references/test-routing.md` for + level-to-command mapping - `review-pr`: provides a structured workflow for reviewing pull requests - `vllm-omni-npu-model-runner-upgrade`: upgrades NPU model runners to align with the latest vllm-ascend NPUModelRunner diff --git a/.claude/skills/vllm-omni-test/SKILL.md b/.claude/skills/vllm-omni-test/SKILL.md new file mode 100644 index 000000000..40a81f4cf --- /dev/null +++ b/.claude/skills/vllm-omni-test/SKILL.md @@ -0,0 +1,890 @@ +--- +name: vllm-omni-test +description: Generate and run tests for vllm-project/vllm-omni with CI-aligned levels and markers; wire new tests into Buildkite (test-ready.yml for L1/L2, test-merge.yml for L3, test-nightly.yml for L4). On completion, always provide copy-paste local and CI-like pytest commands plus prerequisites. Use when creating regression tests, adding L1-L4 coverage, selecting pytest markers, or validating fixes from issues/PRs. +--- + +# vLLM-Omni Test Generator & Runner + +## Purpose + +Use this skill to generate minimal, stable test cases and run them with the correct marker/level strategy for [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni). + +**Link convention:** Paths such as `.buildkite/` and `docs/contributing/` live at the **vllm-omni repo root**. Markdown links use repo-relative paths from this skill file (e.g. `../../../.buildkite/test-ready.yml`, `../../../docs/contributing/ci/CI_5levels.md`). + +Default priorities: + +1. Reproducible regression coverage for bug fixes +2. Correct test level and marker selection +3. Low flake, low dependency tests first +4. CI-compatible run commands +5. **Actionable run commands for the human**: whenever you add or change tests, always finish with **copy-paste-ready** `pytest` lines (local: single file and/or single test; CI-like: markers + `--run-level`), plus short **prerequisites** (GPU tier, HF cache, optional `model_prefix`). Do not assume the reader will infer commands from `test-routing.md` alone. + +## Inputs + +- Issue/PR link and summary +- Changed files or suspected code path +- Whether the user wants local quick validation or CI-equivalent validation +- Hardware constraints (CPU only / CUDA / ROCm / NPU) + +## Workflow + +### Step 1: Classify Test Goal + +- **Bugfix regression**: start from a minimal failing scenario and add assertions that prevent recurrence. Before writing tests, output **`required` / `recommended` / `not_needed`**: **`required`** — stable logic/contract bug that should have been caught; **`recommended`** — environment-sensitive but a small regression still helps; **`not_needed`** — one-off external/config failure or existing tests already cover the path. Prefer the narrowest stable L1 (CPU) case; escalate to L2/L3 only when the bug needs real weights or serving. +- **Feature coverage**: verify new behavior and one negative/boundary case. +- **Perf/benchmark claim**: require benchmark-oriented tests and explicit metrics. + +### Step 2: Select Test Level + +- **L1**: unit/logic, deterministic, CPU-friendly, fastest feedback. +- **L2**: basic e2e/integration and platform-dependent checks. +- **L3/L4**: advanced model/integration/perf validation. + +Use [references/test-routing.md](references/test-routing.md) for level-to-marker and command mapping. + +### Step 3: Pick Markers + +Always attach markers deliberately: + +- **Level**: `core_model` (L1/L2) and/or `advanced_model` (L3) and/or `full_model` (L4 nightly) +- **Model type** (required on model-centric e2e — pick exactly one): + - `omni` — end-to-end multimodal LLM pipelines (thinker/talker/stages; Qwen-Omni family) + - `tts` — speech synthesis / TTS-only models (`/v1/audio/speech`, voice clone, etc.) + - `diffusion` — generative diffusion models (image / audio / text / video from noise) +- **Cross-cutting area** (when relevant): `parallel`, `cache`, `example`, `benchmark` +- **Hardware**: `cpu`, `gpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `distributed_cuda`, … +- Optional: `slow`, distributed markers when multi-card is required + +**Baseline smoke (L2 + L3):** The simplest e2e case per model — default deploy, minimal request — should usually carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** test function so `test-ready.yml` and `test-merge.yml` share one test. `send_*_request` picks validation depth from `--run-level`. References: `test_voxcpm2_tts.py::test_text_to_audio_001`, `test_qwen3_tts_customvoice.py::test_text_to_audio_001`. Heavier scenarios use **`advanced_model` only**; L4 expansion uses **`full_model`**. + +For hardware-aware tests, prefer `@hardware_test(...)` or `hardware_marks(...)` in `tests/helpers/mark.py`. + +**Diffusion nightly split** (`test-nightly.yml`): all diffusion tests use `pytest.mark.diffusion`, but CI groups them by **output modality**, not by separate pytest markers: + +| Nightly group | Scope | Typical models / paths | +|---------------|--------|-------------------------| +| **Diffusion X2I(&A&T) Model Test** | x2**i** (image), x2**a** (audio), x2**t** (text) and other **non-video** diffusion | Qwen-Image*, BAGEL, FLUX, SD3, Z-Image, LongCat, DreamZero, … — `test_*_expansion.py` under X2I shards; L4 sweep uses `-k "not test_wan and not test_bagel_expansion and not hunyuan"` for L4 | +| **Diffusion X2V Model Test** | x2**v** (video) only | Wan2.2, HunyuanVideo 1.5, Wan VACE, … — e.g. `test_wan22_expansion.py`, `test_hunyuan_video_15_expansion.py` | + +Wire new **video** diffusion expansion tests into the **X2V** group; wire **image/audio/text** diffusion expansion tests into **X2I(&A&T)**. Do not place x2v modules in X2I shards (see comment in `test-nightly.yml` above the X2I group). + +### Naming: generated test module files (L2–L4, model-centric e2e) + +When adding **new** pytest modules whose primary scope is a **specific model** (typical under `tests/e2e/offline_inference/` or `tests/e2e/online_serving/`), use this filename pattern: + +| Level | Filename pattern | Example (model `Qwen/Qwen2.5-Omni-7B`) | +|-------|------------------|----------------------------------------| +| **L2**, **L3** | `test_{lowercase_model_slug}.py` | `test_qwen2_5_omni.py` | +| **L4** | `test_{lowercase_model_slug}_expansion.py` | `test_qwen2_5_omni_expansion.py` | + +**Slug rules for `{lowercase_model_slug}`:** + +1. Start from the HuggingFace-style id (e.g. `Qwen/Qwen2.5-Omni-7B`), but **do not** put the org into the filename: use the **repo segment only** (`Qwen2.5-Omni-7B`), not `Qwen_...` / `qwen_qwen2_5_...`. +2. **Lowercase**; replace `.`, `-`, and whitespace with a single `_` (e.g. `Qwen2.5-Omni-7B` → `qwen2_5_omni`). **Omit** trailing size tokens such as `7b` / `30b` in the basename when a single file covers that model line in the directory (matches `test_qwen2_5_omni.py` in-tree). +3. If two checkpoints in the same folder need separate modules, add a **minimal** disambiguator (e.g. `_7b` vs `_3b`) only then. +4. **L1** unit tests are **not** bound to this pattern; use `tests//test_.py` as today. + +Routing tables and commands: [references/test-routing.md](references/test-routing.md) § *Model-centric e2e filename convention*. + +Existing references: `tests/e2e/offline_inference/test_qwen2_5_omni.py` (L2-style omni), `tests/e2e/offline_inference/test_qwen3_5_9b.py` (L2-style omni, single-stage VL), `tests/e2e/online_serving/test_qwen3_omni_expansion.py` (L4-style omni), `tests/e2e/online_serving/test_qwen_image_edit_expansion.py` / `test_qwen_image_expansion.py` (L4-style diffusion). + +### Step 4: Generate Test Case Skeleton + +**1. Pick the functional scenario** (then choose directory, fixtures, and markers): + +| Scenario | Typical location | Fixtures / runner pattern | Baseline markers & level | +|----------|------------------|---------------------------|---------------------------| +| **Offline inference e2e** | `tests/e2e/offline_inference/` | **Module (default):** `omni_runner` + `omni_runner_handler`. **Function (isolation only):** `omni_runner_function` + `omni_runner_handler_function`. Diffusion/TTS may use `Omni(...).generate` directly | L2: `core_model` + **one of** `omni` / `tts` / `diffusion`; `@hardware_test(...)` when GPU/NPU is required | +| **Online serving e2e** | `tests/e2e/online_serving/` | **Module (default):** `omni_server` + `openai_client`. **Function (isolation only):** `omni_server_function` + `openai_client_function`. Clients: `send_omni_request` (omni), `send_audio_speech_request` (tts), `send_diffusion_request` / `send_video_diffusion_request` / `send_images_generations_request` (diffusion) | Baseline smoke: **`core_model` + `advanced_model`**; heavier paths: `advanced_model` only; L4 expansion: `full_model` | +| **Documentation / runnable examples** | `tests/examples/offline_inference/`, `tests/examples/online_serving/` | **Offline docs (preferred):** extract Python/Bash blocks from the doc README (e.g. `ReadmeSnippet.extract_readme_snippets`), `pytest.mark.parametrize` each snippet, run via `example_runner.run` with a stable `output_subfolder`. **Online docs:** copy client/request scripts into dedicated tests and keep them in sync with the doc page. | Usually **L4**: `advanced_model`, often `example` plus hardware marks matching the nightly docs-example job (see `.buildkite/test-nightly.yml`). Full conventions: [docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md) (introduced in [PR #1910](https://github.com/vllm-project/vllm-omni/pull/1910): naming, output directory layout, skip rules, avoid trimming `num_inference_steps` without a strong CI reason). | +| **Performance / benchmark** | `tests/dfx/perf/tests/*.json` + `run_*_benchmark.py` | JSON or script-driven server + load config; assert explicit metrics / baselines | L4 Perf: `full_model` + `benchmark`; wire `test-nightly.yml` Perf steps | +| **Invalid parameter / negative HTTP validation** | `tests/dfx/reliability/invalid_param_test/` | Live `omni_server` + low-level `send_*_http_request` with `err_code` / `err_message` | `pytest.mark.slow` + `omni` / `tts` / `diffusion` + `@hardware_marks` (`H100` or `L4`); CI in **`test-weekly.yml`** (not ready/merge/nightly) | + +**If the user’s test plan includes invalid parameter validation / invalid params / negative HTTP / 400 validation:** do **not** add those `test_*` functions to `tests/e2e/online_serving/test_*.py` or `*_expansion.py`. **Move or author them** under `tests/dfx/reliability/invalid_param_test/` in the **endpoint-matching script** (see **Invalid parameter validation** below). Success-path e2e and invalid-param dfx tests must stay in separate modules. + +**1b. Model type — `omni` vs `tts` vs `diffusion`** + +After choosing offline/online/docs/perf, classify the **product under test** and attach **exactly one** model-type marker. All three can live under the same `tests/e2e/` trees; conventions diverge: + +| Dimension | **Omni** (`pytest.mark.omni`) | **TTS** (`pytest.mark.tts`) | **Diffusion** (`pytest.mark.diffusion`) | +|-----------|------------------------------|----------------------------|----------------------------------------| +| **What it is** | Multimodal LLM pipeline (thinker/talker/stages; text + vision + audio I/O) | Speech synthesis / voice models | Generative diffusion (noise → image, audio, text, or video) | +| **Examples** | Qwen2.5-Omni, Qwen3-Omni | Qwen3-TTS, VoxCPM2, Higgs-Audio, Step-Audio2 | Qwen-Image, BAGEL, Wan2.2, HunyuanVideo | +| **Offline runner** | `omni_runner` + `omni_runner_handler`, `generate_multimodal` | `Omni(...).generate` or stage YAML + TTS params | `Omni(...).generate` + `OmniDiffusionSamplingParams` | +| **Online client** | `openai_client.send_omni_request` (chat completions, modalities) | `openai_client.send_audio_speech_request` (`/v1/audio/speech`) | `send_diffusion_request` (chat/T2I), `send_video_diffusion_request` (`/v1/videos`, X2V), or `send_images_generations_http_request` / `send_images_edits_http_request` (DALL-E routes) | +| **Typical assertions** | Stage outputs, text/audio keywords via `OmniRunnerHandler` / response handler | WAV bytes, stream chunks, speech endpoint contract | Image/video dimensions, `final_output_type`, `assert_diffusion_response` | +| **Stage / deploy YAML** | Per-model omni stage configs (`ci/qwen3_omni_moe.yaml`, …) | `qwen3_tts.yaml`, `voxcpm2.yaml`, … | Often default serve; parallel/offload YAML for heavy DiT | +| **Nightly group** (`test-nightly.yml`) | **Omni Model Test** — `-m "full_model and omni"` | **TTS Model Test** — `-m "full_model and tts"` | **Diffusion X2I(&A&T)** *or* **Diffusion X2V** (see Step 3 table; same `diffusion` marker, different YAML group / file shard) | +| **L4 pressure** | Expansion per modality/model as needed | Expansion + accuracy/perf in TTS group | X2I: merge feature combos per [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); X2V: separate nightly group | + +Do **not** mix fixtures across types (e.g. do not use `omni_runner` layout for a pure diffusion or TTS model without mirroring an in-tree test in that family). + +**Diffusion only — X2I(&A&T) vs X2V (nightly routing, not extra markers):** + +- **X2I(&A&T)**: image / audio / text generation — Qwen-Image*, FLUX, SD3, Z-Image, BAGEL (expansion), LongCat, audio diffusion, etc. +- **X2V**: video generation only — Wan2.2, HunyuanVideo 1.5, Wan VACE, LTX video similarity paths. + +When adding a new `test_*_expansion.py`, place it in the matching **nightly group** step (explicit file list in `test-nightly.yml`), not only by marker expression. + +**2. Use the narrowest deterministic skeleton for the scenario** + +*L1 unit / logic (CPU-first):* + +**Mocking rule (L1 only):** use **pytest** integration — **`mocker`** (`pytest-mock`) or **`monkeypatch`** (built-in). **Do not** import or call **`unittest.mock`** (`patch`, `MagicMock`, `@patch`, `with patch(...)`, etc.) in L1 tests; patches must auto-revert with the test lifecycle. + +```python +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_(mocker): + # Prefer mocker.patch / mocker.spy / mocker.Mock — not unittest.mock.patch + fake_fn = mocker.patch("vllm_omni.some.module.expensive_call", return_value=...) + # Act + # Assert + fake_fn.assert_called_once() + + +def test_(monkeypatch): + # For simple env / attribute substitution without a Mock object + monkeypatch.setenv("SOME_FLAG", "1") + monkeypatch.setattr("vllm_omni.some.module.CONST", 42) + ... +``` + +**Avoid in L1:** + +```python +# BAD — do not use in L1 unit tests +from unittest.mock import patch, MagicMock + +@patch("vllm_omni.some.module.fn") +def test_bad(mock_fn): ... + +def test_also_bad(): + with patch("...") as m: ... +``` + +See **L1 unit test constraints (mocking)** below for the full do/don't list. + +*Offline multimodal e2e — **Omni** (representative):* + +```python +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(...) +@pytest.mark.parametrize("omni_runner", test_params, indirect=True) +def test_(omni_runner, omni_runner_handler) -> None: + request_config = {"prompts": ..., "modalities": [...]} # optional: images, videos, audios + omni_runner_handler.send_omni_request(request_config) +``` + +*Offline generative e2e — **Diffusion** (representative):* + +```python +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(...) +@pytest.mark.parametrize("omni_runner", test_params, indirect=True) +def test_text_to_image_001(omni_runner_handler) -> None: + omni_runner_handler.send_diffusion_request({"prompt": "...", "extra_body": {"num_inference_steps": 4, ...}}) +``` + +*Offline TTS e2e — **Qwen3-TTS** (two-stage; representative):* + +```python +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(...) +@pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True) +def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None: + omni_runner_handler.send_audio_speech_request({ + "input": "...", + "task_type": "Base", + "ref_audio": REF_AUDIO_URL, + "ref_text": REF_TEXT, + }) +``` + +*Offline TTS e2e — **single-stage** (Coqui XTTS, MOSS-TTS-Nano; representative):* + +```python +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(...) +@pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True) +def test_voice_clone_001(omni_runner_handler) -> None: + omni_runner_handler.send_single_stage_tts_request({ + "input": "...", + "language": "en", + "prompt_audio_path": REF_AUDIO_PATH, + "response_format": "wav", + "run_level": "advanced_model", + }) +``` + +*Online serving e2e — **Omni** (representative):* + +```python +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(...) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_to_text_001(omni_server, openai_client) -> None: + request_config = {"model": omni_server.model, "messages": ..., "modalities": ["text"]} + openai_client.send_omni_request(request_config) +``` + +*Online serving e2e — **TTS** (representative):* + +```python +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(...) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_audio_001(omni_server, openai_client) -> None: + request_config = {"model": omni_server.model, "input": "...", "response_format": "wav", ...} + openai_client.send_audio_speech_request(request_config) +``` + +*Online serving e2e — **Diffusion** X2I (representative):* + +```python +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.parametrize("omni_server", _get_default_case(MODEL), indirect=True) +def test_text_to_image_001(omni_server, openai_client) -> None: + openai_client.send_diffusion_request({...}) # chat completions + extra_body +``` + +*Online serving e2e — **Diffusion** X2V (representative):* + +```python +@pytest.mark.core_model +@pytest.mark.diffusion +def test_text_to_video_001(omni_server, openai_client) -> None: + openai_client.send_video_diffusion_request({"model": ..., "form_data": {...}}) # /v1/videos +``` + +*Documentation example tests:* follow the **Preferred Test Strategy** in [l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md): dynamic extraction for offline READMEs; explicit copied client code for online pages until extraction is justified; use the documented **naming**, **output directory** (page folder + case id), and **skipping** rules (e.g. Gradio-only scripts). + +*Performance tests:* add or extend entries under `tests/perf/` (and JSON configs where the project uses them), with **explicit baselines** and the same marker/run-level pairing as the CI step that will execute them. + +**3. Cross-cutting rules** + +- Reuse existing fixtures for the chosen scenario; do not mix “online client” assumptions into offline `OmniRunner` tests without a clear reason. +- Avoid external network dependency in assertions unless the scenario is explicitly “online serving” or doc examples that require a model hub (then align with CI secrets/cache). +- Keep **one test function = one intent** (one modality combo, one endpoint contract, or one acceleration combo). +- **E2E test function layout**: **one case → one `test_` function** with a name that states what is validated (endpoint, size/`n`, server flag, or route). **Do not** merge multiple cases into a single test that branches on `request.node.callspec.id`, `param.id`, or `if case_id == ...`. Use `@pytest.mark.parametrize("omni_server", [...], indirect=True)` per function (usually **one** `OmniServerParams` per test). A loop inside one test is OK only when it serves **that function’s single intent** (e.g. three standard sizes in `test_*_sizes_256_512_1024`). +- **Runtime fixture scope** (`tests/helpers/fixtures/runtime.py`): default **`omni_server` / `omni_runner`** (module) + matching client/handler; use **`omni_server_function` / `omni_runner_function`** only when each `test_*` must start a fresh instance (see below). +- **L1 mocks**: never `unittest.mock`; use `mocker` or `monkeypatch` only (see below). +- **API calls (L2+ e2e, online and offline)**: **reuse** `send_*_request` in `tests/helpers/runtime.py` when it exists; **otherwise add** the helper in `runtime.py` first, then call it from the test. General `assert_*` inside `send_*_request`; special `assert_*` only in the test. See **Runtime send helpers** below — do not call `omni.generate`, raw HTTP, or SDK clients from `test_*.py`. +- **Response assertions**: reusable checks on API bodies / decoded media belong in **`tests/helpers/assertions.py`** — not as private `_assert_*` helpers inside `test_*.py` (see below). +- **Model-specific payloads stay in test modules** — per-model `MODEL`, deploy path, `REF_AUDIO_URL`, `get_prompt()`, `_build_request()`, and inline `request_config` dicts live in `test_{slug}.py` / `test_{slug}_expansion.py` (and offline/L1 siblings). **Do not** create `tests/helpers/{slug}.py` to deduplicate them; a little copy across files is preferred (see `test_glm_tts.py`, `test_cosyvoice3_tts_expansion.py`). `tests/helpers/` is repo-wide harness only (`mark`, `media`, `runtime`, `stage_config`, `assertions`, `fixtures/`). + +#### L1 unit test constraints (mocking) + +L1 tests (`core_model and cpu`, under `tests/diffusion/`, `tests/engine/`, `tests/model_executor/`, etc.) must follow the repo’s **pytest-mock** convention (`pytest-mock>=3.10.0` in `pyproject.toml` `[project.optional-dependencies] dev`). + +| Do | Don't | +|----|-------| +| `def test_foo(mocker):` + `mocker.patch(...)`, `mocker.spy(...)`, `mocker.Mock()`, `mocker.MagicMock()`, `mocker.AsyncMock()` | `from unittest.mock import patch, MagicMock, Mock, AsyncMock` | +| `def test_bar(monkeypatch):` + `monkeypatch.setattr` / `setenv` / `delenv` / `setitem` | `@patch(...)` decorator | +| Let `mocker` auto-stop patches after the test | `with patch(...):` / `patch.object(...)` context managers | +| Mirror neighboring L1 tests in the same directory | `unittest.mock.create_autospec` unless an existing file already documents an exception | + +**Rationale:** `mocker` ties patch lifecycle to pytest fixtures (no leaked patches across tests). `unittest.mock` decorators/context managers are easy to compose incorrectly with parametrized or async tests and are inconsistent with in-tree L1 style. + +**Minimal patterns:** + +```python +def test_returns_cached_config(mocker): + loader = mocker.patch( + "vllm_omni.foo.load_yaml", + return_value={"stages": []}, + ) + result = get_deploy_config("ci/foo.yaml") + assert result["stages"] == [] + loader.assert_called_once_with("ci/foo.yaml") + + +def test_skips_when_env_unset(monkeypatch): + monkeypatch.delenv("VLLM_OMNI_FEATURE", raising=False) + assert should_enable_feature() is False +``` + +E2E levels (L2+) generally **avoid mocks**; if a rare L2 stub is unavoidable, still prefer `mocker` over `unittest.mock` for consistency. + +#### Runtime fixtures — scope (`tests/helpers/fixtures/runtime.py`) + +vLLM-Omni e2e tests start a real **OmniServer** (online) or **OmniRunner** (offline). Pick scope by **how often the process must be recreated**, not by test level alone. + +| Scope | Online fixtures | Offline fixtures | When to use | +|-------|-----------------|------------------|-------------| +| **Module** (default) | `omni_server` → `openai_client` | `omni_runner` → `omni_runner_handler` | **Default for L2/L3/L4 expansion** — amortize model/server init across `test_*` in the same module. Same `OmniServerParams` / runner config can be reused by multiple tests. | +| **Function** | `omni_server_function` → `openai_client_function` | `omni_runner_function` → `omni_runner_handler_function` | **Only when required** — each `test_*` must get a **clean** server/runner (no shared engine/GPU state). Typical: `tests/dfx/reliability/`, sleep/wakeup, crash/restart, tests that mutate global server state. | + +**Rules:** + +1. **Default to module scope** (`omni_server` / `omni_runner`) unless the scenario **explicitly** needs a fresh instance per test function. +2. **Indirect parametrize name must match the fixture name**: `@pytest.mark.parametrize("omni_server", ...)` with `omni_server` + `openai_client`; `@pytest.mark.parametrize("omni_server_function", ...)` with `omni_server_function` + `openai_client_function`. Do not mix module fixture with function client (or vice versa). +3. Different `OmniServerParams` per `test_*` (e.g. default vs `--enable-cpu-offload`) is still OK with **module** `omni_server` — pytest parametrizes per test node; only switch to `_function` when isolation between tests matters, not merely because `server_args` differ. +4. One `test_*` with **many** server configs (expansion matrix) → single function + `@pytest.mark.parametrize("omni_server", [...], indirect=True)` + module `omni_server` (see in-tree `test_qwen_image_expansion.py`). + +```python +# Default — module-scoped server (L4 expansion) +@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=H100)], indirect=True) +def test_foo_images_generations_default_1024(omni_server, openai_client) -> None: + openai_client.send_images_generations_request({...}) + +# Function-scoped — reliability / per-test clean state only +@pytest.mark.parametrize( + "omni_server_function", + [pytest.param(OmniServerParams(model=MODEL), marks=H100)], + indirect=True, +) +def test_foo_sleep_wakeup_cycle(omni_server_function, openai_client_function) -> None: + openai_client_function.send_omni_sleep_http_request({...}) + openai_client_function.send_omni_wakeup_http_request({...}) +``` + +#### Runtime send helpers — online **and** offline (`tests/helpers/runtime.py`) + +**L2+ e2e (online serving and offline inference) must call APIs through `tests/helpers/runtime.py`.** Fixtures live in `tests/helpers/fixtures/runtime.py`; **send/assert implementation** lives in `tests/helpers/runtime.py` + `tests/helpers/assertions.py`. + +| Principle | Action | +|-----------|--------| +| **Reuse first** | Grep `runtime.py` for an existing **`send_*_request`** on `OpenAIClientHandler` (online) or `OmniRunnerHandler` (offline) that matches the endpoint / pipeline shape. | +| **Extend when close** | If an existing `send_*_request` almost fits (e.g. missing one optional field), extend it in `runtime.py` — do not fork logic in the test file. | +| **Add when missing** | No suitable helper → add **`send__request`** (high-level: call + general `assert_*`) or **`send___http_request`** (low-level HTTP for negative/dfx) in **`runtime.py` first**, then call it from tests. | +| **Test module owns payload only** | In `test_*.py`: `MODEL`, deploy path, vendored media, `get_prompt()`, and inline **`request_config` dicts** only. **No** `omni.generate(...)`, raw `requests.post`, OpenAI SDK calls, or `_collect_audio()` / `_process_output()` in e2e tests. | + +**Online** (`openai_client` from `omni_server`): `OpenAIClientHandler.send_*_request`. + +**Offline** (`omni_runner_handler` from `omni_runner`): `OmniRunnerHandler.send_*_request`. + +| Do | Don't | +|----|-------| +| Online: `openai_client.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_video_diffusion_request`, `send_images_generations_request`, … | `requests.post(f"{base_url}/v1/...", json=…)` or `client.chat.completions.create(...)` inside a test | +| Offline: `omni_runner_handler.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_single_stage_tts_request`, `send_single_stage_tts_batch_request`, … | `omni_runner.omni.generate(...)` + hand-rolled tensor/WAV extraction in `test_*.py` | +| Add missing `send_*` to **`runtime.py`** first; bundle general `assert_*` inside the send helper | A one-off `def _post_*` / `def _collect_audio` at the bottom of a test module | +| Mirror naming/style of neighboring `send_*` (docstring, `request_config` dict, optional `run_level`, `err_code` / `err_message` for negative cases) | Different parameter shapes per test file for the same endpoint | + +**Workflow when generating online/offline e2e tests:** + +1. Decide whether the needed check is **general** (every success call of this `send_*`) or **special** (one case / one parameter combo only). See **General vs special assert placement** below. +2. **Search `runtime.py`** for a matching **`send_*_request`**; if missing, add it (with general `assert_*` inside) before writing the test body. +3. In the test module: build `request_config` → call **`send_*_request` only** for the general contract; call **`assert_*` in the test file only** for special, case-specific checks (import from `assertions.py`). +4. Reserve **low-level** `send_*_http_request` for negative/dfx tests (`err_code` / `err_message`) — not for ordinary L2+ success-path e2e. +5. When wiring Buildkite `source_file_dependencies`, include **`tests/helpers/runtime.py`** and/or **`tests/helpers/assertions.py`** when new helpers were added. + +**Exceptions (document in the test docstring why):** models whose offline prompt path cannot go through existing handlers yet (e.g. `test_higgs_audio_v2.py`, `test_voxtral_tts.py` with custom tokenizer compose) may call `omni.generate` directly until a `send_*_request` is added to `runtime.py` — treat as **debt**, not the default for new TTS/diffusion/omni e2e. + +**General vs special assert placement:** + +| Kind | Definition | Where it lives | Test file calls | +|------|------------|----------------|-----------------| +| **General** | Default success contract for this endpoint/client method — e.g. HTTP 200, `data[]` shape, `n` count, `size` dimensions, decodable image/audio, bundled omni/diffusion fields | Implement in `assertions.py`, invoke **inside** `send_*_request` in `runtime.py` (after low-level send / SDK call, when `err_code` is not set) | **`send_*_request` only** — do **not** repeat the same `assert_*` | +| **Special** | Extra check tied to **this test case only** — e.g. seed byte-identical replay, accuracy/CLIP threshold, perf ceiling, model-specific optional field | Add or reuse `assert_*` in `assertions.py` (never inline in test) | `send_*_request` **then** `assert_(...)` **once** for that case | + +Changing `request_config` fields (`size`, `n`, `seed`, server flags) is **not** special validation — the general `assert_*` should read those from `request_config` inside `send_*`. + +**`send_*` ↔ `assert_*` pairing (online `OpenAIClientHandler`):** + +| High-level `send_*` (prefer in L2+ e2e) | Assert already invoked inside `runtime.py` | Low-level HTTP-only sibling | +|----------------------------------------|--------------------------------------------|-----------------------------| +| `send_omni_request` | `assert_omni_response` | `send_chat_completions_http_request` → `assert_http_error` only | +| `send_diffusion_request` | `assert_diffusion_response` | — | +| `send_audio_speech_request` | `assert_audio_speech_response` | `send_audio_speech_http_request` → `assert_http_error` only | +| `send_video_diffusion_request` | `assert_diffusion_response` | `send_videos_*_http_request` → `assert_http_error` only | +| `send_images_generations_request` *(add when needed)* | **`assert_images_generations_response`** *(general — inside send)* | `send_images_generations_http_request` → `assert_http_error` only | +| `send_images_edits_request` *(add when needed)* | **`assert_images_edits_response`** *(general — inside send)* | `send_images_edits_http_request` → `assert_http_error` only | + +**Rule:** Tests call high-level `send_*_request` for the general contract. **Never** call the same bundled `assert_*` again. Call an extra `assert_*` in the test **only** for special, case-specific validation. + +**Common `OpenAIClientHandler` entry points** (non-exhaustive — grep `runtime.py` before adding): + +| Area | High-level (SDK + assert) | Low-level HTTP (`*_http_request`) | +|------|---------------------------|-----------------------------------| +| Omni chat | `send_omni_request` | `send_chat_completions_http_request` | +| Diffusion T2I (chat route) | `send_diffusion_request` | — | +| Diffusion X2V | `send_video_diffusion_request` | `send_videos_create_http_request`, `send_video_content_http_request`, … | +| DALL-E T2I / edit | `send_images_generations_request`, `send_images_edits_request` | `send_images_generations_http_request`, `send_images_edits_http_request` | +| TTS (online) | `send_audio_speech_request` | `send_audio_speech_http_request`, `send_audio_generate_http_request`, … | +| Ops / meta | — | `send_health_http_request`, `send_models_http_request`, `send_omni_sleep_http_request`, … | + +**Common `OmniRunnerHandler` entry points** (offline — grep `runtime.py` before adding): + +| Area | High-level (`send_*_request` + assert) | Notes | +|------|----------------------------------------|-------| +| Omni multimodal | `send_omni_request` | `generate_multimodal` path | +| Diffusion offline | `send_diffusion_request` | chat-route / `OmniTextPrompt` offline | +| Qwen-style TTS | `send_audio_speech_request` | two-stage, `generate_multimodal` + `mm_processor_kwargs` | +| Single-stage TTS (Coqui XTTS, MOSS-TTS-Nano, …) | `send_single_stage_tts_request`, `send_single_stage_tts_batch_request` | `prompt` + `additional_information` + `omni.generate` — **do not** duplicate in tests | +| New model family | **Add** `send__request` here first | Then call from `tests/e2e/offline_inference/test_*.py` | + +L1 tests under `tests/entrypoints/` may use **FastAPI `TestClient`** or direct handler calls with mocks; they do **not** need `OpenAIClientHandler` / `OmniRunnerHandler`. **L2+ online and offline e2e** must use `runtime.py` `send_*` helpers. + +#### Invalid parameter validation (`tests/dfx/reliability/invalid_param_test/`) + +When the user asks for **invalid parameter validation**, **invalid request bodies**, **HTTP 4xx contract tests**, or any case that sends **malformed / out-of-range / mismatched** API payloads against a **live** server: + +1. **Do not** place these in `tests/e2e/online_serving/` or `*_expansion.py`. If already drafted there, **move** the `test_*` into the correct `invalid_param_test` script and delete the duplicate from e2e. +2. **Pick the script by HTTP route** (extend an existing file; add a new `test_invalid_.py` only when no in-tree script covers that route family): + +| API / area | Script | +|------------|--------| +| Omni chat completions, WebSocket video/realtime paths | `test_invalid_omni_chat.py` | +| `POST /v1/audio/speech`, stream, batch, voices | `test_invalid_audio_speech.py` | +| Audio diffusion endpoints | `test_invalid_audio_diffusion.py` | +| `POST /v1/images/generations` | `test_invalid_image_generation.py` | +| `POST /v1/images/edits` | `test_invalid_image_editing.py` | +| `POST/GET/DELETE /v1/videos*` | `test_invalid_video_generation.py` | +| Sleep / wakeup / server control | `test_invalid_server_control.py` | + +3. **Match in-tree style** in the chosen script: + +| Element | Convention | +|---------|------------| +| **Module markers** | `pytestmark = [pytest.mark.slow, pytest.mark.]` | +| **Server fixture** | `_PARAMS` / `_QWEN3_TTS_SPEECH`-style list of `pytest.param(OmniServerParams(...), id="...", marks=hardware_marks(...))`; `@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)` | +| **Hardware** | `hardware_marks(res={"cuda": "H100"})` for heavy diffusion/omni/video; `hardware_marks(res={"cuda": "L4"})` for smaller TTS models (must match weekly `-m "slow and L4"` step) | +| **HTTP client** | **Low-level** `openai_client.send_*_http_request({..., "err_code": 400, "err_message": (...)} )` — **not** `send_*_request` (success path) | +| **Case shape** | Prefer **one `test_*` per route family** + `@pytest.mark.parametrize("body_spec, err_message", [...])` with stable `id=` per case; or dedicated `test__malformed_json` when not parametrized | +| **Body helpers** | `_minimal__json(omni_server)` / `_minimal_*_form_data()`; merge overrides with `body.update(body_spec)` | +| **Known gaps** | `pytest.mark.skip(reason="…#3649")` as `_SKIP_ISSUE_3649` when server validation is not yet strict (mirror neighboring cases) | +| **Sections** | Route banner comments (`# ─── POST /v1/images/generations ───`) like existing files | +| **Shared fixtures** | Reuse `tests/dfx/reliability/invalid_param_test/conftest.py` (`tiny_png_bytes`, env defaults) | + +4. **Adding a new model** to invalid-param coverage: append a `pytest.param(OmniServerParams(model="...", stage_config_path=..., server_args=...), ...)` entry to the script’s `_PARAMS` list — do **not** create `test_invalid_.py` unless the route is new. + +5. **CI — `.buildkite/test-weekly.yml` only** (not `test-ready.yml` / `test-merge.yml` / `test-nightly.yml`): + +| Weekly step | Command | When your cases run | +|-------------|---------|---------------------| +| **Invalid parameters Test · H100** | `pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100"` | Diffusion / omni / video invalid-param tests with `H100` hardware mark | +| **Invalid parameters Test · L4** | `pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4"` | TTS / lighter models marked `L4` | + +- **Trigger:** `build.env("WEEKLY") == "1"` or PR label `weekly-test`. +- **Default:** extending an existing `invalid_param_test` script needs **no YAML edit** — the weekly steps already sweep the whole directory. +- **Edit YAML only** when adding a **new hardware queue**, a **new top-level script** that must run in isolation, or a **model-specific weekly shard** (mirror neighboring reliability steps). +- Weekly steps **do not** use `source_file_dependencies`. + +**Example — append to `test_invalid_image_generation.py`:** + +```python +@pytest.mark.parametrize( + "body_spec, err_message", + [ + pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"), + ], +) +@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True) +def test_images_generations_invalid_requests( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + body_spec: dict[str, object], + err_message: str | tuple[str, ...], +) -> None: + body = _minimal_images_gen_json(omni_server) + body.update(body_spec) + openai_client.send_images_generations_http_request( + {"json": body, "timeout": 300, "err_code": 400, "err_message": err_message} + ) +``` + +See [references/test-routing.md](references/test-routing.md) **Invalid parameter / weekly CI**. + +#### Assertion helpers (`tests/helpers/assertions.py`) + +**Do not** add module-local helpers such as `_assert_images_generations_payload` or `_send_and_assert_*` in e2e test files. Response/media validation belongs in **`tests/helpers/assertions.py`**, grouped by **category**: + +| Category | Existing anchors | When to extend vs add | +|----------|------------------|------------------------| +| **Image** (chat diffusion + DALL-E JSON) | `assert_image_diffusion_response`, `assert_image_valid` | DALL-E `/v1/images/generations` JSON → add **`assert_images_generations_response`** beside image helpers (reuses `assert_image_valid`); do not fork decode logic into tests | +| **Video** | `assert_video_diffusion_response`, `assert_video_valid` | Extend these for new video contracts | +| **Audio** | `assert_audio_diffusion_response`, `assert_audio_speech_response`, `assert_audio_valid` | Extend for new TTS/audio endpoints | +| **Omni multimodal** | `assert_omni_response` | Extend for new modality combos | +| **HTTP errors** | `assert_http_error`, `assert_err_message_in_text` | Used inside low-level `send_*_http_request` and negative tests | + +| Do | Don't | +|----|-------| +| Put **general** `assert_*` inside `send_*_request` in `runtime.py`; tests call `send_*_request` only | Call `assert_diffusion_response` after `send_diffusion_request` (general assert already bundled) | +| Put **special** `assert_*` in `assertions.py` and call it **in the test** after `send_*_request` when that case needs extra checks | Put general decode/count/size logic in the test because “this case uses `n=4`” (that belongs in general assert reading `request_config`) | +| Extend the matching category function, or add `assert__response` next to its category | `_send_and_assert_*` or PIL/base64 loops in `test_*.py` | +| Use low-level `send_*_http_request` in dfx/negative tests (`err_code`) | Use low-level HTTP send + manual general assert in every L2+ expansion file | + +**Workflow when generating tests:** + +1. Implement or reuse **general** `assert_*` in `assertions.py` (by category). +2. Wire it into **`send_*_request`** in `runtime.py` if not already bundled. +3. In the test: `request_config` → `send_*_request`. Add a separate `assert_*` import/call **only** for special case validation. +4. Update `source_file_dependencies` when `runtime.py` or `assertions.py` changes. + +**Example — DALL-E `/v1/images/generations`:** + +```python +# tests/helpers/assertions.py — general contract (image category) +def assert_images_generations_response(resp_body: dict, request_config: dict, *, run_level: str | None = None) -> None: + ... # data[], n, size → assert_image_valid + +# tests/helpers/runtime.py — general assert INSIDE send +def send_images_generations_request(self, request_config: dict[str, Any], ...) -> list[HttpResponse]: + responses = self.send_images_generations_http_request(request_config, ...) + if request_config.get("err_code") is None: + body = responses[0].json_body + assert isinstance(body, dict) + assert_images_generations_response(body, request_config, run_level=self.run_level) + return responses + +# tests/e2e/.../test_foo_expansion.py — one case per test_* (no case_id branching) +@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True) +def test_foo_images_generations_default_1024(omni_server, openai_client) -> None: + openai_client.send_images_generations_request({"json": body, "timeout": 300}) + +@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True) +def test_foo_images_generations_sizes_256_512_1024(omni_server, openai_client) -> None: + for size in ("256x256", "512x512", "1024x1024"): + openai_client.send_images_generations_request({"json": {**body, "size": size}, "timeout": 300}) + +# chat diffusion — separate test; send_diffusion_request bundles assert_diffusion_response +@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True) +def test_foo_chat_completions_t2i_fallback(omni_server, openai_client) -> None: + openai_client.send_diffusion_request({...}) +``` + +L1 tests under `tests/entrypoints/` may keep **minimal** asserts next to the handler under test when validating a single branch; still prefer `assertions.py` when the same JSON/media check appears in more than one file or level (e.g. L1 protocol test + L4 e2e). + +#### Omni Test Writing Guidance (L1-L4 Layering) + +When the goal is a general Omni/multimodal test case, prioritize mapping the test to the correct purpose, directory, and resource assumptions aligned with the layers (see [CI_5levels.md](../../../docs/contributing/ci/CI_5levels.md)): + +- **L1**: Unit/logic validation on CPU (`core_model and cpu`). Cover input validation, branches, and exception paths (`tests//test_*.py`). **Mock with `mocker` / `monkeypatch` only — not `unittest.mock`.** +- **L2**: Basic e2e (online/offline basic scenarios). Prefer dummy/lightweight models to validate the end-to-end request-to-output-structure/streaming chain (typically `tests/e2e/online_serving/` and `tests/e2e/offline_inference/`). +- **L3/L4**: Important integration, performance, and accuracy validation. L4 emphasizes “full functional scenarios + performance/stress + runnable doc examples” (typically `*_expansion.py` plus related expansion cases). + +Also keep markers consistent with the run level: use `core_model` for L1/L2 and `advanced_model` for L3/L4, and pair with `--run-level` to select the intended CI strategy. + +#### Diffusion Test Writing Guidance (L4 Coverage Combinations) + +When the task involves diffusion models/features, organize L4 test cases following [`#1832`](https://github.com/vllm-project/vllm-omni/issues/1832): combine multiple diffusion features into as few test cases as possible to fit limited CI GPU resources. + +Implementation strategy: + +1. If “full L4 is too heavy”: first provide a **reduced local validation case in L1/L2** (so the key assertion and the contract fix point are covered deterministically). +2. Then provide **CI/nightly-ready L3/L4 high-marked cases** (e.g. `advanced_model`) to broaden coverage under resource constraints. + +#### L4 nightly sub-pillars — Function vs Accuracy vs Perf + +**L4 is not only `test_*_expansion.py`.** In `test-nightly.yml`, each model-type group typically runs **separate jobs**: + +| Sub-pillar | Nightly step label pattern | What you add | Typical paths | +|------------|----------------------------|--------------|---------------| +| **Function** | `· Function Test with …` | E2e expansion / feature matrix | `tests/e2e/online_serving/test__expansion.py` (or offline) | +| **Accuracy** | `· Accuracy Test` | Quality / similarity vs baseline | `tests/e2e/accuracy/test_*.py` | +| **Perf** | `· Perf Test · ` | Throughput / latency / memory benchmarks | `tests/dfx/perf/tests/test__vllm_omni.json` + runner script | +| **Doc** (optional) | `· Doc Test` | Runnable doc examples | `tests/examples/*/test_text_to_image.py`, … | + +**Default when the user asks for “L4 functional cases”:** deliver **Function** pillar only (`*_expansion.py` + Function Test shard in `test-nightly.yml`). **Do not** silently add Perf or Accuracy unless the user also asks for **performance** / **benchmark** / **accuracy** / **full L4** / **full L4 coverage**. + +**When the user explicitly asks for L4 perf (or full L4 including perf):** + +1. **Do not** put throughput/latency assertions inside `test_*_expansion.py` — perf uses the **dfx benchmark harness**, not `omni_server` e2e fixtures. +2. Add or extend a **JSON config** under `tests/dfx/perf/tests/` (mirror in-tree names: `test_qwen_image_vllm_omni.json`, `test_bagel_vllm_omni.json`). +3. Each JSON entry: `test_name`, `server_params` (`model`, `serve_args`), `benchmark_params[]` with explicit **`baseline`** (`throughput_qps`, `latency_mean`, `peak_memory_mb_mean`, …). +4. Run locally via the matching script (model type → runner): + +| Model type | Runner script | Example config | +|------------|---------------|----------------| +| **Diffusion X2I/X2V** | `tests/dfx/perf/scripts/run_diffusion_benchmark.py` | `tests/dfx/perf/tests/test_qwen_image_vllm_omni.json` | +| **TTS** | `tests/dfx/perf/scripts/run_benchmark.py` | `tests/dfx/perf/tests/test_tts.json` (shared) or `test_{slug}.json` (dedicated; use when the model must not join the shared nightly matrix before integration — e.g. `test_voxcpm2.json`, `test_coqui_tts.json`) | +| **Omni** | (see existing Omni Perf steps — `run_benchmark.py` / dedicated omni configs in-tree) | `tests/dfx/perf/tests/test_omni*.json` | + +5. Wire **`test-nightly.yml` · Perf Test** step for that model (separate from Function Test): export `DIFFUSION_BENCHMARK_DIR` / `BENCHMARK_DIR`, run pytest on the script with `--test-config-file`, **upload artifacts** (`buildkite-agent artifact upload`), often **multi-GPU H100** for diffusion perf. +6. **One benchmark scenario → one `test_name` block** in JSON (same spirit as one `test_*` per function case). Combine server `serve_args` + workload in one entry; use multiple `benchmark_params` rows for size/step sweeps under the same server config. + +**Perf local command (diffusion example):** + +```bash +cd tests +export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results +export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN +pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py \ + --test-config-file dfx/perf/tests/test__vllm_omni.json +``` + +**Clarify in the test plan** which L4 pillars you deliver: `Function only` | `Function + Perf` | `Function + Accuracy + Perf`. + +See [references/test-routing.md](references/test-routing.md) **L4 nightly pillars** for Buildkite step patterns. + +### Step 5: Wire Buildkite (when CI must run the new test) + +If the test is not already collected by an existing pipeline command (for example, L1 tests marked `core_model and cpu` are already covered by the **Simple Unit Test** step in ready/merge), update the appropriate pipeline under `.buildkite/`: + +| Test level | Edit this file | Typical trigger / intent | +|------------|----------------|---------------------------| +| **L1** and **L2** | [`.buildkite/test-ready.yml`](../../../.buildkite/test-ready.yml) | PR **ready** label; L1 CPU + L2 GPU/basic e2e (steps labeled **Omni ·**, **TTS ·**, **Diffusion ·**) | +| **L3** | [`.buildkite/test-merge.yml`](../../../.buildkite/test-merge.yml) | Post-merge; `advanced_model` integration per model type | +| **L4** | [`.buildkite/test-nightly.yml`](../../../.buildkite/test-nightly.yml) | Nightly; grouped by model type (see below) | +| **Invalid param / reliability (weekly)** | [`.buildkite/test-weekly.yml`](../../../.buildkite/test-weekly.yml) | Weekly; `tests/dfx/reliability/invalid_param_test/` — **not** L1–L4 e2e pipelines | + +**Level-specific Buildkite delivery (follow the user’s requested level):** + +| User asks for | Document / edit **only** | Do **not** add by default | +|---------------|---------------------------|---------------------------| +| **L4** (`*_expansion.py`, `full_model`) | `test-nightly.yml` — append file to the matching nightly shard (X2I / X2V / Omni / TTS) **or** note it is already collected by an existing `-m` / `-k` sweep | `test-merge.yml` (L3 / `advanced_model`) or `test-ready.yml` (L2) | +| **L3** | `test-merge.yml` with `advanced_model` + `source_file_dependencies` | `test-nightly.yml` unless user also wants nightly | +| **L2** | `test-ready.yml` with `core_model` + `source_file_dependencies` | merge / nightly pipelines | +| **Invalid param** | `test-weekly.yml` — **Invalid parameters Test** group (`-m "slow and H100"` / `-m "slow and L4"`). Usually **no YAML edit** when appending cases to existing `invalid_param_test` scripts | `test-ready.yml`, `test-merge.yml`, `test-nightly.yml` | + +`test-nightly.yml` shards use **explicit pytest file paths** in `commands` (and PR labels like `diffusion-x2iat-test`); they generally **do not** use `source_file_dependencies`. Only suggest merge/ready YAML when the requested level is L3/L2 (or the user explicitly asks for multi-level CI). + +**`test-nightly.yml` top-level groups** (each uses `full_model` + type marker in pytest `-m`): + +| Group | Type marker | Notes | +|-------|-------------|--------| +| **Omni Model Test** | `omni` | Function / doc / accuracy / **perf** for Qwen-Omni family | +| **TTS Model Test** | `tts` | `-m "full_model and L4 and tts"` (and H100 variants when enabled); **Perf Test** uses `run_benchmark.py` + `test_tts.json` | +| **Diffusion X2I(&A&T) Model Test** | `diffusion` | Image/audio/text diffusion; explicit file shards for **Function**; **Perf Test · <Model>** uses `run_diffusion_benchmark.py` + JSON config | +| **Diffusion X2V Model Test** | `diffusion` | Video-only (Wan, HunyuanVideo, …); PR label `diffusion-x2v-test`; separate **Perf** steps when in-tree | + +Guidelines when editing YAML: + +- **Match markers and `--run-level`** to the test: L1/L2 in `test-ready.yml` use `core_model` + `omni` / `tts` / `diffusion` as appropriate; L3 merge uses `advanced_model`; L4 nightly uses `full_model` with the same type marker. +- **Diffusion L4**: add x2i/x2a/x2t expansion files to an **X2I(&A&T)** step; add x2v expansion files to **X2V** — never rely on a broad `tests/e2e/` sweep alone (shards exist to save GPU and avoid wrong queue). +- **`source_file_dependencies` (required for new E2E jobs in `test-ready.yml` and `test-merge.yml`)**: every step under the **`:card_index_dividers: E2E Test`** group **must** declare `source_file_dependencies` so Buildkite only schedules the job when relevant paths change. Omitting it breaks path-based CI skipping for GPU e2e jobs. +- **Reuse or extend an existing step** when the new test shares the same marker expression, queue, and timeout; otherwise add a new `steps:` entry with the correct `agents.queue`, `timeout_in_minutes`, and docker/kubernetes plugin blocks consistent with neighboring jobs. +- **Job timeout — use `timeout_in_minutes` on the step, not `timeout` before pytest**: when adding or documenting a new E2E job in `test-ready.yml` / `test-merge.yml`, set `timeout_in_minutes` on the step and run `pytest` directly in `commands`. **Do not** wrap pytest in `timeout 40m bash -c "..."` / `timeout 20m bash -c "..."`; Buildkite already enforces the step deadline via `timeout_in_minutes`. +- **Never omit `plugins` in skill examples or PR YAML snippets**: H100 E2E steps on `mithril-h100-pool` use the full `kubernetes` block (`resources`, `volumeMounts`, `env` with `HF_TOKEN` secret, `nodeSelector: gpu-h100-sxm`, `devshm` + `hf-cache` volumes). L4/docker steps on `gpu_1_queue` use the full `docker#v5.2.0` block (`shm-size`, `HF_HOME`, volumes). Do not write `# ... plugins unchanged` — paste the complete block from a sibling step. +- **Platform forks** (e.g. AMD-ready / AMD-merge) live alongside these files; apply the same level → file mapping for those pipelines when the test is platform-specific. + +#### `source_file_dependencies` for E2E Test jobs (`test-ready.yml`, `test-merge.yml`) + +When adding or splitting a step inside the **E2E Test** group, **always** include `source_file_dependencies` on that step. List paths the job actually depends on — at minimum: + +1. **Test module(s)** — every `tests/e2e/...` file the `pytest` command runs (online and offline if both are in one step). +2. **Model implementation** — `vllm_omni/model_executor/models//` and/or `vllm_omni/diffusion/models//` as applicable. +3. **Stage / input plumbing** (omni & multi-stage TTS) — `vllm_omni/model_executor/stage_input_processors/.py` when present. +4. **Deploy / stage YAML** — `vllm_omni/deploy/.yaml` or `vllm_omni/deploy/ci/.yaml` referenced by the test’s `OmniServerParams` / `get_deploy_config_path(...)`. +5. **Test helpers** (when changed in the same PR) — `tests/helpers/runtime.py` for new `send_*` client methods; `tests/helpers/assertions.py` for new shared `assert_*` helpers. + +Use **repo-relative paths** (no leading `./`). Prefer **directory** entries (`.../models/bagel/`) when the whole package matters; use **file** entries for single test modules and YAML configs. Mirror a neighboring step for the same model family (e.g. **Diffusion · Bagel Test** in `test-ready.yml`). + +**Example** (L2 ready — diffusion online smoke, H100 kubernetes): + +```yaml +- label: "Diffusion · Bagel Test" + source_file_dependencies: + - tests/e2e/online_serving/test_bagel.py + - vllm_omni/model_executor/models/bagel/ + - vllm_omni/diffusion/models/bagel/ + - vllm_omni/model_executor/stage_input_processors/bagel.py + - vllm_omni/deploy/bagel.yaml + timeout_in_minutes: 40 + commands: + - | + export VLLM_IMAGE_FETCH_TIMEOUT=60 + pytest -s -v tests/e2e/online_serving/test_bagel.py -m 'core_model' --run-level 'core_model' + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate +``` + +**Example** (L3 merge — TTS online + offline, L4 docker): + +```yaml +- label: "TTS · Qwen3-TTS Base Test" + source_file_dependencies: + - tests/e2e/online_serving/test_qwen3_tts_base.py + - tests/e2e/offline_inference/test_qwen3_tts_base.py + - vllm_omni/model_executor/models/qwen3_tts/ + - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py + - vllm_omni/deploy/qwen3_tts.yaml + timeout_in_minutes: 20 + commands: + - | + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 + pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model' + agents: + queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" +``` + +**Example** (L3 merge — diffusion online + offline, H100 kubernetes): + +```yaml +- label: "Diffusion · Qwen Image Test" + source_file_dependencies: + - tests/e2e/online_serving/test_qwen_image.py + - tests/e2e/offline_inference/test_qwen_image.py + - vllm_omni/diffusion/models/qwen_image/ + timeout_in_minutes: 40 + commands: + - | + export VLLM_IMAGE_FETCH_TIMEOUT=60 + pytest -s -v \ + tests/e2e/online_serving/test_qwen_image.py \ + tests/e2e/offline_inference/test_qwen_image.py \ + -m 'advanced_model and cuda' --run-level 'advanced_model' + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate +``` + +When extending an **existing** E2E step to run an additional test file, **append** that file (and any new model/deploy paths) to `source_file_dependencies` on the same step. When authoring a PR, mention the dependency list in the **Buildkite change** section of your test plan. + +#### L3 / L4: validating Buildkite from a feature branch (root `pipeline.yml`) + +Production behavior is driven by [`.buildkite/pipeline.yml`](../../../.buildkite/pipeline.yml): it builds the CI image, then uploads **L2** (`test-ready.yml` on non-`main`), **L3** (`test-merge.yml` when `build.branch == "main"`), and **L4** (`test-nightly.yml` when `build.env("NIGHTLY") == "1"`). + +When you add or debug **L3** or **L4** tests and need those child pipelines to run **off `main`** or **without** `NIGHTLY=1`, apply **temporary** edits (revert before merge): + +1. **Point the upload at the right file (optional shortcut)** + In the **Upload Ready Pipeline** step, the command is `buildkite-agent pipeline upload .buildkite/test-ready.yml`. For a one-off run you may change that path to `.buildkite/test-merge.yml` (L3) or `.buildkite/test-nightly.yml` (L4), and adjust the step’s `if:` so it does not conflict with the other upload steps (avoid double-uploading unless intentional). + +2. **L3 — merge pipeline** + The gate `if: build.branch == "main"` lives on the **Upload Merge Pipeline** step in `pipeline.yml` (not inside `test-merge.yml`). **Comment out that `if` line** so `test-merge.yml` is uploaded on your feature branch. Alternatively rely on step (1) instead of the dedicated merge upload step. + +3. **L4 — nightly pipeline** + **Comment out** `if: build.env("NIGHTLY") == "1"` on the **Upload Nightly Pipeline** step in `pipeline.yml` so the nightly definition is uploaded without setting the env var. + Child steps in [`test-nightly.yml`](../../../.buildkite/test-nightly.yml) often repeat `if: build.env("NIGHTLY") == "1"`; **comment out those lines on the steps you need to run**, otherwise they will still be skipped after upload. + +### Step 6: Run Tests + +Pick one command path: + +- **Quick local regression (preferred first)**: single file or `file.py::test_name` +- **CI-like level run**: marker expression + `--run-level` + +Full templates live in [references/test-routing.md](references/test-routing.md). **After authoring tests, always emit concrete commands** (see **Output Format**). + +**Typical copy-paste examples** (run from repo `tests/` directory; requires matching vLLM/vllm-omni install and hardware): + +| Area | Example | +|------|---------| +| **Omni offline L2** | `pytest -s -v e2e/offline_inference/test_qwen2_5_omni.py -m "core_model and omni and not cpu" --run-level=core_model` | +| **Omni online L2** | `pytest -s -v e2e/online_serving/test_qwen3_omni.py -m "core_model and omni" --run-level=core_model` | +| **TTS online L2** | `pytest -s -v e2e/online_serving/test_qwen3_tts_base.py -m "core_model and tts" --run-level=core_model` | +| **Diffusion X2I L2 online** | `pytest -s -v e2e/online_serving/test_qwen_image.py -m "core_model and diffusion" --run-level=core_model` | +| **Diffusion X2V L2 online** | `pytest -s -v e2e/online_serving/test_wan22_t2v.py -m "core_model and diffusion" --run-level=core_model` | +| **Diffusion X2I L4 nightly** | `pytest -s -v e2e/online_serving/test_qwen_image_expansion.py -m "full_model and diffusion and H100" --run-level=full_model` | +| **Diffusion X2I L4 perf (nightly)** | `pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file dfx/perf/tests/test_qwen_image_vllm_omni.json` | +| **Diffusion X2V L4 nightly** | `pytest -s -v e2e/online_serving/test_wan22_expansion.py -m "full_model and cuda" --run-level=full_model` | +| **Invalid param (weekly H100)** | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and H100"` | +| **Invalid param (weekly L4)** | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and L4"` | +| **L1 CPU** | `pytest -s -v -m "core_model and cpu"` | + +**Prerequisites to mention when relevant**: GPU model (e.g. L4 vs H100), `HF_HOME` / token for hub weights, module-level `skipif` (NPU/XPU-only gaps), and whether CI already collects the path (e.g. `test_*_expansion.py` glob in `test-nightly.yml`). + +### Step 7: Validate Result Quality + +Before finishing: + +- Is the new assertion directly tied to the bug/feature contract? +- Are API calls made via **`runtime.py` `send_*_request`** (not inline HTTP/SDK in the test file)? +- Is the **general** `assert_*` bundled inside `send_*_request` (not duplicated in the test)? +- Is each scenario a **separate `test_*` function** (no `case_id` / `if param.id` branching across cases)? +- **Fixture scope**: default **`omni_server` / `omni_runner`** (module); **`_function`** variants only when each test must restart the instance? +- Is the test deterministic (no fragile timing/network coupling)? +- Is runtime appropriate for the selected level? +- Are markers and `--run-level` consistent? +- **Invalid-param cases** in `tests/dfx/reliability/invalid_param_test/` (not e2e), using `send_*_http_request` + `err_code`, with `pytest.mark.slow` and correct `H100` / `L4` hardware mark? + +## Output Format + +When completing a request, return: + +1. **Test plan** (level, markers, file target, **L4 pillar(s)** if applicable: Function / Accuracy / Perf / Doc; **or** **Invalid param** pillar with target `invalid_param_test/test_invalid_.py`; module basename for e2e: `test_{slug}.py` / `test_{slug}_expansion.py`) +2. **Generated/updated test file(s)** — model-specific constants and request payloads stay **inside** those test modules (**never** `tests/helpers/{slug}.py`). Extend **`tests/helpers/runtime.py`** / **`tests/helpers/assertions.py`** only when new **repo-wide** `send_*` or shared assert helpers are required (do not leave ad-hoc HTTP or `_assert_*` in test modules). **If L4 Perf was requested**, also list **`tests/dfx/perf/tests/*.json`**. **If invalid-param cases were requested**, list the **`invalid_param_test/` script** and parametrized `test_*` / new `id=` rows — **not** e2e paths. +3. **Buildkite change** — **match the requested level and pillar**: + - **L4 Function**: `test-nightly.yml` **Function Test** shard (file list / note existing sweep). **No** `test-merge.yml` unless user also asked for L3. + - **L4 Perf**: `test-nightly.yml` **Perf Test · <Model>** step (new step or extend commands); include env exports + artifact upload pattern from a sibling perf job. + - **L4 Accuracy**: `test-nightly.yml` **Accuracy Test** step under the same model-type group. + - **L3**: `test-merge.yml` + full `source_file_dependencies` + `agents` + `plugins`. + - **L2**: `test-ready.yml` + same E2E block requirements. + - **Invalid param**: `test-weekly.yml` — note **Invalid parameters Test · H100/L4** group; usually **no YAML change** when only extending existing scripts. + If the user asked only for **L4 functional cases**, state explicitly that **Perf / Accuracy / Invalid param** were not included (offer to add if needed). If the plan mixes **success e2e** and **invalid param**, split deliverables across e2e vs `invalid_param_test/` and call out both CI files. +4. **Run commands (required)** — always include, in fenced `bash` blocks: + - **Local — whole file**: `cd tests` then `pytest -s -v …` + - **Local — single test** (optional but preferred when the change is one function): `pytest -s -v path::test_func …` + - **CI-like** (when not L1 CPU): the same **marker + `--run-level`** pairing the level uses (see [references/test-routing.md](references/test-routing.md) and **Step 6** table above) + - **Prerequisites** (one line): e.g. “Requires CUDA L4 + weights cached locally”; “Requires H100 + `HF_TOKEN`”; “Syntax-only validation when vllm is not installed locally” +5. **Result summary** (pass/fail if you executed tests; if not executed, state that explicitly and what the user should run) + +## Additional Resources + +- Marker and command routing: [references/test-routing.md](references/test-routing.md) +- CI pipelines (vllm-omni): [test-ready.yml](../../../.buildkite/test-ready.yml) (L1/L2), [test-merge.yml](../../../.buildkite/test-merge.yml) (L3), [test-nightly.yml](../../../.buildkite/test-nightly.yml) (L4), [test-weekly.yml](../../../.buildkite/test-weekly.yml) (invalid param / reliability) +- L4 documentation example tests (naming, extraction vs copied scripts, output dirs, skips): [docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md) — see also [PR #1910](https://github.com/vllm-project/vllm-omni/pull/1910) diff --git a/.claude/skills/vllm-omni-test/references/test-routing.md b/.claude/skills/vllm-omni-test/references/test-routing.md new file mode 100644 index 000000000..445a78399 --- /dev/null +++ b/.claude/skills/vllm-omni-test/references/test-routing.md @@ -0,0 +1,372 @@ +# Test Routing Reference + +Use this reference to map testing goals to levels, markers, and runnable commands. + +**Repo paths** (`.buildkite/`, `docs/contributing/…`): link with repo-relative paths from this file (e.g. `../../../../.buildkite/test-ready.yml`, `../../../../docs/contributing/ci/CI_5levels.md`). + +## Model-centric e2e filename convention (L2–L4) + +When generating a **new** test module tied to a **specific model** under `tests/e2e/offline_inference/` or `tests/e2e/online_serving/`: + +| Level | Pattern | Example (`Qwen/Qwen2.5-Omni-7B`) | +|-------|---------|----------------------------------| +| **L2**, **L3** | `test_{lowercase_model_slug}.py` | `test_qwen2_5_omni.py` | +| **L4** | `test_{lowercase_model_slug}_expansion.py` | `test_qwen2_5_omni_expansion.py` | + +**Slug rules for `{lowercase_model_slug}`** (canonical: [SKILL.md](../SKILL.md) § *Naming: generated test module files*): + +1. Start from the HuggingFace-style id (e.g. `Qwen/Qwen2.5-Omni-7B`), but **do not** put the org into the filename: use the **repo segment only** (`Qwen2.5-Omni-7B`), not `Qwen_...` / `qwen_qwen2_5_...`. +2. **Lowercase**; replace `.`, `-`, and whitespace with a single `_` (e.g. `Qwen2.5-Omni-7B` → `qwen2_5_omni`). **Omit** trailing size tokens such as `7b` / `30b` in the basename when a single file covers that model line in the directory (matches `test_qwen2_5_omni.py` in-tree). +3. If two checkpoints in the same folder need separate modules, add a **minimal** disambiguator (e.g. `_7b` vs `_3b`) **only then**. +4. **L1** unit tests are **not** bound to this pattern; use `tests//test_.py` as today. + +## Model type markers (`omni` / `tts` / `diffusion`) + +Every model-centric e2e test must declare **exactly one** type marker: + +| Marker | Model family | Typical APIs | +|--------|--------------|--------------| +| `pytest.mark.omni` | Multimodal LLM (Qwen-Omni, …) | `send_omni_request`, `omni_runner`, stage YAML | +| `pytest.mark.tts` | TTS / speech synthesis | `send_audio_speech_request`, `/v1/audio/speech` | +| `pytest.mark.diffusion` | Diffusion generative models | `send_diffusion_request`, `send_video_diffusion_request`, `send_images_generations_http_request`, `OmniDiffusionSamplingParams` | + +**Diffusion** uses one marker for all output modalities. Nightly CI splits diffusion by **group**, not by extra markers: + +| Nightly group (`test-nightly.yml`) | Diffusion scope | Examples | +|-----------------------------------|-----------------|----------| +| **Diffusion X2I(&A&T) Model Test** | x2**i** / x2**a** / x2**t** — image, audio, text (non-video) | Qwen-Image*, BAGEL, FLUX, SD3, Z-Image, LongCat, DreamZero | +| **Diffusion X2V Model Test** | x2**v** — video only | Wan2.2, HunyuanVideo 1.5, Wan VACE | + +PR labels for selective nightly runs: `diffusion-x2iat-test`, `diffusion-x2v-test` (plus `nightly-test` / `NIGHTLY=1`). + +## Level and Marker Mapping + +| Goal | Suggested Level | Marker baseline | Typical location | +|------|------------------|-----------------|------------------| +| Unit logic, regression on pure code path | L1 | `core_model and cpu` | `tests//test_*.py` | + +### L1 unit tests — mocking (`mocker`, not `unittest.mock`) + +L1 modules run in **Simple Test** (`test-ready.yml` / `test-merge.yml`: `-m 'core_model and cpu'`). When isolation requires doubles: + +- **Use** the **`mocker`** fixture (`pytest-mock`): `mocker.patch`, `mocker.spy`, `mocker.Mock`, `mocker.MagicMock`, `mocker.AsyncMock`. +- **Use** **`monkeypatch`** for env vars and simple `setattr` / `delattr` without mock objects. +- **Do not use** `unittest.mock` — no `from unittest.mock import ...`, no `@patch`, no `with patch(...)`. + +```python +# Good (L1) +def test_handler(mocker): + mocker.patch("vllm_omni.pkg.fn", return_value=0) + +# Bad (L1) +from unittest.mock import patch + +@patch("vllm_omni.pkg.fn") +def test_handler(mock_fn): ... +``` + +E2E (L2+) should not rely on mocks unless documenting a rare exception; prefer real (or lightweight) execution paths. +| Basic integration/e2e | L2 | `core_model` + **one of** `omni` / `tts` / `diffusion` | `tests/e2e/...` | +| Advanced integration | L3 | `advanced_model` + type marker | `tests/e2e/...` | +| Full function / nightly | L4 | `full_model` (nightly) + type marker | **Function:** `tests/e2e/*_expansion.py`; **Perf:** `tests/dfx/perf/tests/*.json`; **Accuracy:** `tests/e2e/accuracy/` | +| Invalid HTTP / param validation | Weekly (dfx) | `pytest.mark.slow` + type marker + `H100` or `L4` | `tests/dfx/reliability/invalid_param_test/test_invalid_*.py` | + +## Marker Selection Rules + +1. **Level** (pick one): + - `core_model` — L1/L2 (`test-ready.yml`) + - `advanced_model` — L3 (`test-merge.yml`) + - `full_model` — L4 nightly (`test-nightly.yml`); some expansion tests still carry both `advanced_model` and `full_model` during migration +2. **Model type** (pick one): `omni`, `tts`, or `diffusion` +3. **Cross-cutting** when relevant: `parallel`, `cache`, `example`, `benchmark` +4. **Hardware**: `cpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `distributed_cuda`, … +5. Multi-card: `@hardware_test(...)` in `tests/helpers/mark.py` + +## Command Templates + +### Quick local checks + +```bash +cd tests +pytest -s -v test_xxxx.py +``` + +### L1 + +```bash +cd tests +pytest -s -v -m "core_model and cpu" +``` + +### L2 + +```bash +cd tests +pytest -s -v -m "core_model and not cpu" --run-level=core_model +``` + +### L3 (merge) + +```bash +cd tests +pytest -s -v -m "advanced_model" --run-level=advanced_model +``` + +### L4 (nightly) + +**Function** (e2e expansion — default when user asks for “L4 functional cases”): + +```bash +cd tests +pytest -s -v e2e/online_serving/test_qwen_image_expansion.py -m "full_model and diffusion and H100" --run-level=full_model +``` + +**Perf** (dfx benchmark harness — separate pillar; not `omni_server` fixtures): + +```bash +cd tests +export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results +export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN +pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py \ + --test-config-file dfx/perf/tests/test_qwen_image_vllm_omni.json +``` + +Broad marker sweep (when no explicit file shard): + +```bash +cd tests +pytest -s -v -m "full_model" --run-level=full_model +``` + +## L4 nightly pillars (Function / Accuracy / Perf) + +Nightly **L4** for a model is often **multiple jobs** in `test-nightly.yml`, not a single pytest module: + +| Pillar | User says | Deliver | Nightly step | +|--------|-----------|---------|--------------| +| **Function** | “L4 functional cases” / “functional” / `*_expansion.py` | `tests/e2e/.../test__expansion.py` | `· Function Test with H100/L4` | +| **Perf** | “performance” / “perf” / “benchmark” / “full L4” | `tests/dfx/perf/tests/test__vllm_omni.json` + runner | `· Perf Test · ` | +| **Accuracy** | “accuracy” / “similarity” | `tests/e2e/accuracy/test_*.py` | `· Accuracy Test` | + +**Rule:** “L4 functional cases” → **Function only** by default. State that Perf/Accuracy are separate pillars; add them only when requested. + +**Diffusion perf JSON** (`tests/dfx/perf/tests/test__vllm_omni.json`): array of objects with `test_name`, `server_params` (`model`, `serve_args`), `benchmark_params[]`, and per-workload **`baseline`** (`throughput_qps`, `latency_mean`, `peak_memory_mb_mean`). Copy structure from `test_qwen_image_vllm_omni.json` in-tree. + +**Diffusion perf nightly step** (under **Diffusion X2I(&A&T)**; copy full `kubernetes` plugins from `· Perf Test · Qwen-Image`): + +```yaml + - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · " + key: nightly-diffusion-x2iat-performance- + timeout_in_minutes: 180 + commands: + - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results + - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN + - export CACHE_DIT_VERSION=1.3.0 + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test__vllm_omni.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" + buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" + exit $$EXIT + agents: + queue: "mithril-h100-pool" + plugins: + # … full kubernetes block (often multi-GPU) — copy from Perf Test · Qwen-Image +``` + +| Model type | Perf runner | Config example | +|------------|-------------|----------------| +| Diffusion X2I/X2V | `dfx/perf/scripts/run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json` | +| TTS | `dfx/perf/scripts/run_benchmark.py` | `test_tts.json` (`BENCHMARK_DIR`) | +| Omni | `run_benchmark.py` / omni configs in-tree | `test_omni*.json` | + +Do **not** put throughput/latency baselines inside `test_*_expansion.py` — that belongs in the dfx perf JSON + nightly Perf job. + +## Invalid parameter validation (weekly / dfx) + +When the user’s test plan includes **invalid parameter validation**, **invalid request bodies**, or **HTTP 4xx** validation against a live server: + +1. **Do not** author these in `tests/e2e/online_serving/test_*.py` or `*_expansion.py`. **Move** any drafted `test_*` into `tests/dfx/reliability/invalid_param_test/`. +2. **Pick script by route** (extend in-tree file): + +| Route family | Script | +|--------------|--------| +| Omni chat / WS video·realtime | `test_invalid_omni_chat.py` | +| `/v1/audio/speech` (+ stream / batch / voices) | `test_invalid_audio_speech.py` | +| Audio diffusion | `test_invalid_audio_diffusion.py` | +| `/v1/images/generations` | `test_invalid_image_generation.py` | +| `/v1/images/edits` | `test_invalid_image_editing.py` | +| `/v1/videos*` | `test_invalid_video_generation.py` | +| Sleep / wakeup / server control | `test_invalid_server_control.py` | + +3. **Style:** `pytestmark = [pytest.mark.slow, pytest.mark.]`; `_PARAMS` + `hardware_marks`; `send_*_http_request` with `err_code` + `err_message`; parametrized `body_spec` rows with `id=`; `_minimal_*_json()` helpers; `_SKIP_ISSUE_3649` when tracked in [#3649](https://github.com/vllm-project/vllm-omni/issues/3649). +4. **CI:** [`.buildkite/test-weekly.yml`](../../../../.buildkite/test-weekly.yml) group **Reliability Test - Invalid parameters Test** — **not** ready/merge/nightly. + +```bash +# Weekly H100 (diffusion / omni / video invalid-param) +cd tests +pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and H100" + +# Weekly L4 (e.g. Qwen3-TTS speech invalid-param) +pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and L4" + +# Single script / case +pytest -s -v dfx/reliability/invalid_param_test/test_invalid_image_generation.py::test_images_generations_invalid_requests -m "slow and H100" +``` + +**Trigger:** `WEEKLY=1` or PR label `weekly-test`. Appending cases to existing scripts usually needs **no YAML edit** (directory sweep). No `source_file_dependencies` on weekly steps. + +### Platform-targeted examples + +```bash +cd tests +pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model +``` + +### Concrete e2e paths (common in-tree) + +Paths are relative to `tests/` after `cd tests`. + +| Scenario | Example command | +|----------|------------------| +| **Omni** offline L2 | `pytest -s -v e2e/offline_inference/test_qwen2_5_omni.py -m "core_model and omni and not cpu" --run-level=core_model` | +| **Omni** offline L2 — one test | `pytest -s -v e2e/offline_inference/test_qwen2_5_omni.py::test_text_to_text -m "core_model and omni and not cpu" --run-level=core_model` | +| **Omni** online L2 | `pytest -s -v e2e/online_serving/test_qwen3_omni.py -m "core_model and omni" --run-level=core_model` | +| **Omni** offline L2 (Qwen3.5-9B VL) | `pytest -s -v e2e/offline_inference/test_qwen3_5_9b.py -m "core_model and omni and not cpu" --run-level=core_model` | +| **TTS** online L2 | `pytest -s -v e2e/online_serving/test_qwen3_tts_base.py -m "core_model and tts" --run-level=core_model` | +| **TTS** online L2 — one test | `pytest -s -v e2e/online_serving/test_qwen3_tts_base.py::test_text_to_audio_001 -m "core_model and tts" --run-level=core_model` | +| **Diffusion X2I** offline L2 | `pytest -s -v e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion and not cpu" --run-level=core_model` | +| **Diffusion X2I** L2 online smoke | `pytest -s -v e2e/online_serving/test_qwen_image.py e2e/online_serving/test_bagel.py -m "core_model and diffusion" --run-level=core_model` | +| **Diffusion X2V** L2 online smoke | `pytest -s -v e2e/online_serving/test_wan22_t2v.py -m "core_model and diffusion" --run-level=core_model` | +| **Diffusion X2I** L4 online expansion (H100) | `pytest -s -v e2e/online_serving/test_qwen_image_expansion.py -m "full_model and diffusion and H100" --run-level=full_model` | +| **Diffusion X2I** L4 perf (nightly) | `pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file dfx/perf/tests/test_qwen_image_vllm_omni.json` | +| **Diffusion X2I** L4 online expansion (Edit) | `pytest -s -v e2e/online_serving/test_qwen_image_edit_expansion.py -m "full_model and diffusion and H100" --run-level=full_model` | +| **Diffusion X2V** L4 nightly | `pytest -s -v e2e/online_serving/test_wan22_expansion.py -m "full_model and cuda" --run-level=full_model` | +| **Invalid param** weekly H100 | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and H100"` | +| **Invalid param** weekly L4 | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and L4"` | + +### Agent / author completion checklist + +When adding or modifying tests, do not stop at “where the file lives” — also deliver: + +1. **Local**: `cd tests` + `pytest -s -v ` (and `path::test_func` when a single case is enough). +2. **CI-like**: marker string + `--run-level` matching the pipeline (`core_model` / `advanced_model` / `full_model`). +3. **L1 mocks**: `mocker` / `monkeypatch` only; no `unittest.mock`. +4. **API client + assert placement**: **General** validation → implement in `assertions.py`, call **inside** `send_*_request` in `runtime.py`; tests call **`send_*_request` only**. **Special** case validation → `assert_*` in `assertions.py`, called **in the test** after `send_*_request`. Low-level `send_*_http_request` is for negative/dfx tests (`err_code`), not ordinary L2+ success e2e. +5. **Shared assertions**: logic in **`tests/helpers/assertions.py`**; general checks inside `send_*_request`; special checks only in the test. No `_assert_*` in `test_*.py`. +6. **One case → one `test_*`**: function name reflects what is validated; no `if case_id == ...` mega-test merging multiple scenarios. +7. **Fixture scope**: default **`omni_server` + `openai_client`** / **`omni_runner` + `omni_runner_handler`** (module). Use **`omni_server_function` + `openai_client_function`** / **`omni_runner_function` + `omni_runner_handler_function`** only when each `test_*` must spawn a fresh instance. Parametrize name must match fixture (`omni_server` vs `omni_server_function`). +8. **Type marker**: `omni`, `tts`, or `diffusion` on every model e2e module. +9. **Diffusion L4 Function**: wire `*_expansion.py` into **X2I(&A&T)** or **X2V** **Function Test** in **`test-nightly.yml` only** — do not add `test-merge.yml` unless the user also requested L3. +10. **Diffusion L4 Perf** (only when requested): add `tests/dfx/perf/tests/test__vllm_omni.json` + **Perf Test · <Model>** step (artifact upload); not part of “L4 functional cases” by default. +11. **E2E Buildkite (L2/L3 only)**: `test-ready.yml` / `test-merge.yml` steps need **`source_file_dependencies`** + full **`agents` + `plugins`**. **L4 nightly** uses explicit file lists or perf scripts in `test-nightly.yml` (no merge job). +12. **Invalid param**: cases in **`tests/dfx/reliability/invalid_param_test/`** (route-matching script), `send_*_http_request` + `err_code`, `pytest.mark.slow` + `H100`/`L4`; CI = **`test-weekly.yml`** only — do not put in e2e or nightly. +13. **Prerequisites**: GPU tier, HF cache/token, and any module `skipif` / platform-only YAML. + +## Buildkite pipeline mapping + +| Level | Repo file | Model-type grouping | +|-------|-----------|---------------------| +| L1, L2 | [`.buildkite/test-ready.yml`](../../../../.buildkite/test-ready.yml) | Steps prefixed **Omni ·**, **TTS ·**, **Diffusion ·** under **E2E Test** — **`source_file_dependencies` required** | +| L3 | [`.buildkite/test-merge.yml`](../../../../.buildkite/test-merge.yml) | Per-model E2E steps; `-m "advanced_model and …"` — **`source_file_dependencies` required** | +| L4 | [`.buildkite/test-nightly.yml`](../../../../.buildkite/test-nightly.yml) | **Omni / TTS / Diffusion X2I(&A&T) / X2V** — each group may have **Function**, **Accuracy**, **Perf**, **Doc** steps | +| Invalid param (weekly) | [`.buildkite/test-weekly.yml`](../../../../.buildkite/test-weekly.yml) | **Invalid parameters Test · H100** / **· L4** — sweeps `tests/dfx/reliability/invalid_param_test/` | + +### `source_file_dependencies` (E2E Test only — ready & merge) + +Required on each step inside the **E2E Test** group in `test-ready.yml` and `test-merge.yml`. Typical entries: + +| Category | Path pattern | +|----------|----------------| +| Test module | `tests/e2e/online_serving/test_.py`, `tests/e2e/offline_inference/test_.py` | +| API client helpers | `tests/helpers/runtime.py` (when the e2e job uses newly added `send_*` methods) | +| Shared assert helpers | `tests/helpers/assertions.py` (when the e2e job uses newly added `assert_*` helpers) | +| AR / omni model | `vllm_omni/model_executor/models//` | +| Diffusion model | `vllm_omni/diffusion/models//` | +| Stage processor | `vllm_omni/model_executor/stage_input_processors/.py` | +| Deploy config | `vllm_omni/deploy/.yaml` or `vllm_omni/deploy/ci/.yaml` | + +Copy the dependency block from the closest in-tree E2E step for the same model family. Update the list whenever the pytest command or server YAML changes. + +### E2E `agents` + `plugins` blocks (required — do not truncate) + +When generating or documenting E2E steps, **always include the complete `plugins` section**. Set **`timeout_in_minutes`** on the step for the job deadline; run **`pytest` directly** in `commands` (no `timeout 40m bash -c` wrapper around pytest). + +Common patterns in `test-ready.yml` / `test-merge.yml`: + +| Queue | Plugin | Required fields | +|-------|--------|-----------------| +| `mithril-h100-pool` | `kubernetes` | `resources.limits.nvidia.com/gpu`, `volumeMounts` (`devshm`, `hf-cache`), `env` (`HF_HOME`, `HF_TOKEN` via `secretKeyRef`), `nodeSelector: gpu-h100-sxm`, `volumes` (`devshm` emptyDir, `hf-cache` hostPath `/mnt/hf-cache`) | +| `gpu_1_queue` / `gpu_4_queue` | `docker#v5.2.0` | `image`, `always-pull`, `propagate-environment`, `shm-size: "8gb"`, `environment: HF_HOME`, `volumes: /fsx/hf_cache` (add `HF_TOKEN` when the step needs hub access) | + +**H100 kubernetes template** (copy from **Diffusion · Bagel Test** / **Diffusion · Qwen Image Test** in-tree): + +```yaml + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate +``` + +**L4 docker template** (copy from **TTS · Qwen3-TTS Base Test** in `test-merge.yml`): + +```yaml + agents: + queue: "gpu_1_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" +``` + +When extending an existing step (e.g. add `tests/e2e/offline_inference/test_qwen_image.py` to **Diffusion · Qwen Image Test**), update `source_file_dependencies` and `commands` only; **keep the existing `agents` + `plugins` block unchanged** unless the hardware tier changes. + +The root [`.buildkite/pipeline.yml`](../../../../.buildkite/pipeline.yml) decides **which** child file is uploaded. To run **L3** on a feature branch, comment out `if: build.branch == "main"` on the merge upload step. To run **L4** without `NIGHTLY=1`, comment out the nightly upload step’s `if: build.env("NIGHTLY") == "1"` and the same `if` on relevant steps inside `test-nightly.yml`. Revert such edits before merging. + +Platform-specific pipelines (e.g. AMD) follow the same level → file pairing under `.buildkite/`. + +## Diffusion RFC (#1832) Alignment Tips + +For **X2I(&A&T)** coverage planning: + +- Prioritize high-value feature combinations with minimal case count. +- Split into lightweight L1/L2 validation plus L4 `*_expansion.py` under the **X2I** nightly shards. +- **X2V** models (Wan, HunyuanVideo) stay in the **X2V** group — do not merge into X2I sweeps. + +If hardware is insufficient, provide an executable reduced case plus a deferred full CI/nightly plan. diff --git a/docs/contributing/ci/CI_5levels.md b/docs/contributing/ci/CI_5levels.md index 01310606e..adf29bfa5 100644 --- a/docs/contributing/ci/CI_5levels.md +++ b/docs/contributing/ci/CI_5levels.md @@ -11,6 +11,10 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari Level Scope & Focus + Model Coverage Strategy + Feature Coverage Strategy + Interface Coverage Strategy + Tags Time Cost Test Dir Doc @@ -24,6 +28,10 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari Contribution Guideline & PR checklist / / + / + / + / + / .github/PULL_REQUEST_TEMPLATE.md Test Style (PR Checklist) / / @@ -32,6 +40,10 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari CI Failure Description / / + / + / + / + / CI Failures / / @@ -39,6 +51,10 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari L1
(Unit & Logic) Unit tests for components like entrypoints, models + / + / + / + core_model and cpu <15min /tests/{component_name}/test_xxx @@ -50,10 +66,18 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari L2
(E2E across models & GPU-required UT) - Online & Offline (basic deployment scenarios):
dummy, normal inference function (output format, stream), some instance startup UT + Online (basic deployment scenarios):
dummy, normal inference function (output format, stream), some instance startup UT + High-priority models + online basic scenarios + request success validation + High-priority features (using random lightweight models) + High-priority interfaces (using random lightweight models) + core_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion + Model tests:
/tests/e2e/online_serving/test_{model_name}.py
- /tests/e2e/offline_inference/test_{model_name}.py + Feature tests:
+ /tests/{component_name}/test_xxx
+ Interface tests:
+ /tests/entrypoints/test_xxx Chapter 1
@@ -65,10 +89,19 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari L3
(Important Perf & Integration & Accuracy) Online & Offline (multiple deployment scenarios):
real model, normal inference function, normal accuracy + High/medium-priority models with real weights + online/offline key scenarios + basic accuracy validation + Medium-priority features (using random lightweight models) + Medium-priority interfaces (using random lightweight models) + advanced_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion <30min + Model tests:
/tests/e2e/online_serving/test_{model_name}.py
- /tests/e2e/offline_inference/test_{model_name}.py + /tests/e2e/offline_inference/test_{model_name}.py
+ Feature tests:
+ /tests/{component_name}/test_xxx
+ Interface tests:
+ /tests/entrypoints/test_xxx Chapter 2
@@ -79,18 +112,27 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari L4
(Perf & Integration & Accuracy) - Online & Offline: full functional scenarios + performance test + doc test + Online: full functional scenarios + performance test + doc test + accuracy test + High-priority models: function, performance, accuracy, and doc testing
Medium/low-priority models: function and doc testing + Low-priority features (using random lightweight models) + Low-priority interfaces (using random lightweight models) + full_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion <3 hour - Full Function:
+ Model tests:
/tests/e2e/online_serving/test_{model_name}_expansion.py
- /tests/e2e/offline_inference/test_{model_name}_expansion.py
+ Feature tests:
+ /tests/{component_name}/test_xxx
+ Interface tests:
+ /tests/entrypoints/test_xxx
Performance:
/tests/dfx/perf/tests/test_qwen_omni.json (Omni), test_tts.json (TTS),
and /tests/dfx/perf/tests/test_{diffusion_model}_vllm_omni.json (Diffusion)
Doc Test:
tests/example/online_serving/test_{model_name}.py
- tests/example/offline_inference/test_{model_name}.py + tests/example/offline_inference/test_{model_name}.py
+ Accuracy Test:
+ /tests/e2e/accuracy/test_{model_name}.py Chapter 3
@@ -101,7 +143,11 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari L5
(Stability & Reliability) - Online & Offline: long-term stability test + reliability test + Online: long-term stability test + reliability test + Long-term stability and reliability testing for high-priority models + / + Invalid-parameter validation for high-priority interfaces + slow and hardware_test(H100, L4, etc.) and omni/tts/diffusion Depends on reality Stability:
From f75c437a60884691c0da88cbbaec9e1ba27bcec4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B1=AA=E5=BF=97=E9=B9=8F?= Date: Thu, 18 Jun 2026 17:14:18 +0800 Subject: [PATCH 05/40] [Refactor]Base class for dit pipelines with unified parameter declaration[1/N] (#4225) Signed-off-by: princepride Co-authored-by: Cursor Co-authored-by: Hongsheng Liu Co-authored-by: Claude Sonnet 4.6 --- examples/offline_inference/bagel/README.md | 391 ----------------- examples/offline_inference/bagel/end2end.py | 278 ------------ .../image_to_image/image_edit.py | 100 +++-- .../text_to_image/text_to_image.py | 125 ++++-- examples/online_serving/bagel/README.md | 404 ------------------ .../bagel/openai_chat_client.py | 184 -------- examples/online_serving/bagel/run_server.sh | 12 - .../bagel/run_server_stage_cli.sh | 164 ------- .../online_serving/image_to_image/README.md | 23 +- .../image_to_image/openai_chat_client.py | 44 +- .../sensenova_u1/openai_chat_client.py | 4 +- recipes/Bagel/BAGEL-7B-MoT.md | 227 ++++++++++ recipes/README.md | 1 + recipes/SenseNova/SenseNova-U1.md | 25 ++ tests/model_extras/test_model_extras.py | 157 +++++++ vllm_omni/diffusion/data.py | 61 +++ vllm_omni/diffusion/diffusion_engine.py | 4 - vllm_omni/diffusion/io_support.py | 28 -- .../sensenova_u1/pipeline_sensenova_u1.py | 24 -- vllm_omni/diffusion/utils/param_utils.py | 27 ++ vllm_omni/engine/async_omni_engine.py | 16 + vllm_omni/entrypoints/openai/serving_chat.py | 51 +-- vllm_omni/model_extras/__init__.py | 20 + vllm_omni/model_extras/bagel.py | 74 ++++ vllm_omni/model_extras/registry.py | 147 +++++++ vllm_omni/model_extras/sensenova_u1.py | 21 + 26 files changed, 1026 insertions(+), 1586 deletions(-) delete mode 100644 examples/offline_inference/bagel/README.md delete mode 100644 examples/offline_inference/bagel/end2end.py delete mode 100644 examples/online_serving/bagel/README.md delete mode 100755 examples/online_serving/bagel/openai_chat_client.py delete mode 100755 examples/online_serving/bagel/run_server.sh delete mode 100644 examples/online_serving/bagel/run_server_stage_cli.sh create mode 100644 recipes/Bagel/BAGEL-7B-MoT.md create mode 100644 tests/model_extras/test_model_extras.py create mode 100644 vllm_omni/diffusion/utils/param_utils.py create mode 100644 vllm_omni/model_extras/__init__.py create mode 100644 vllm_omni/model_extras/bagel.py create mode 100644 vllm_omni/model_extras/registry.py create mode 100644 vllm_omni/model_extras/sensenova_u1.py diff --git a/examples/offline_inference/bagel/README.md b/examples/offline_inference/bagel/README.md deleted file mode 100644 index f2685d56f..000000000 --- a/examples/offline_inference/bagel/README.md +++ /dev/null @@ -1,391 +0,0 @@ -# BAGEL-7B-MoT - -## Setup - -Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. - -## Architecture - -BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies: - -| Topology | Stages | Description | -| :------- | :----- | :---------- | -| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. | -| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. | - -Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`. - -## Quick Start - -```bash -cd examples/offline_inference/bagel - -# Default two-stage mode (auto-detected) -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A cute cat" - -# Single-stage mode -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A cute cat" \ - --deploy-config vllm_omni/deploy/bagel_single_stage.yaml -``` - -> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices. - -## Modality Control - -Control the mode using the `--modality` argument: - -| Modality | Input | Output | Description | -| :------- | :---- | :----- | :---------- | -| `text2img` | Text | Image | Generate images from text prompts | -| `img2img` | Image + Text | Image | Transform images using text guidance | -| `img2text` | Image + Text | Text | Generate text descriptions from images | -| `text2text` | Text | Text | Pure text generation (language model mode) | - -### Text to Image (text2img) - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A cute cat" \ - --steps 50 -``` - -### Image to Image (img2img) - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality img2img \ - --image-path /path/to/image.jpg \ - --prompts "Let the woman wear a blue dress" -``` - -### Image to Text (img2text) - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality img2text \ - --image-path /path/to/image.jpg \ - --prompts "Describe this image in detail" -``` - -### Text to Text (text2text) - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2text \ - --prompts "What is the capital of France?" - -# Load prompts from a text file (one prompt per line): -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2text \ - --txt-prompts /path/to/prompts.txt -``` - -## Think Mode - -Think mode enables the model to generate `...` planning/reasoning tokens before producing the final output. This improves generation quality for complex prompts. - -- **Two-stage**: The Thinker (AR) stage decodes think tokens, then transfers the augmented KV cache to the DiT stage for image generation. -- **Single-stage**: The DiT's internal LLM generates think tokens in-place before proceeding to denoise. - -```bash -# Think + text2img: plan before generating -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A futuristic city with flying cars" \ - --think \ - --max-think-tokens 1000 - -# Think + img2img: reason about the edit -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality img2img \ - --image-path /path/to/image.jpg \ - --prompts "Make it look like a watercolor painting" \ - --think - -# Think + img2text: reason before describing -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality img2text \ - --image-path /path/to/image.jpg \ - --prompts "What is happening in this image?" \ - --think - -# Think + text2text: chain-of-thought reasoning -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2text \ - --prompts "Solve: 23 * 47" \ - --think -``` - -Think mode parameters: - -| Argument | Default | Description | -| :------- | :------ | :---------- | -| `--think` | `False` | Enable thinking mode | -| `--max-think-tokens` | `1000` | Maximum tokens for think generation | -| `--do-sample` | `False` | Enable sampling (vs. greedy) for text generation | -| `--text-temperature` | `0.3` | Temperature for text generation sampling | - -## Classifier-Free Guidance (CFG) - -CFG controls the trade-off between prompt fidelity and diversity. These parameters apply to image generation modalities (`text2img`, `img2img`). - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A photorealistic portrait" \ - --cfg-text-scale 6.0 \ - --cfg-img-scale 2.0 \ - --negative-prompt "blurry, low quality, distorted" \ - --cfg-interval 0.4 1.0 \ - --cfg-renorm-type global \ - --cfg-renorm-min 0.0 -``` - -| Argument | Default | Description | -| :------- | :------ | :---------- | -| `--cfg-text-scale` | `4.0` | Text CFG scale (higher = more prompt-adherent) | -| `--cfg-img-scale` | `1.5` | Image CFG scale (for img2img) | -| `--negative-prompt` | `None` | Negative prompt for CFG conditioning | -| `--cfg-interval` | pipeline default | CFG active interval `[start, end]` as fractions of total timesteps | -| `--cfg-renorm-type` | `None` | Renormalization type: `global`, `text_channel`, `channel` | -| `--cfg-renorm-min` | `None` | Minimum renormalization value | -| `--cfg-parallel-size` | `1` | CFG parallel size: `1` = batched (single GPU), `2` = 2-branch parallel, `3` = full 3-GPU parallel | - -## Deployment Topologies - -### Two-Stage (Default) - -The default topology auto-detected from the model. No extra flags needed. - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A cute cat" -``` - -The pipeline is defined in [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml). Stage 0 (Thinker) and Stage 1 (DiT) share GPU 0 by default. For dual-GPU setups, customize the deploy YAML and set `devices: "1"` for stage 1. - -### Single-Stage - -Pass the single-stage deploy config via `--deploy-config`: - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A cute cat" \ - --deploy-config vllm_omni/deploy/bagel_single_stage.yaml -``` - -See [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) for configuration details. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry. - -### Tensor Parallelism (TP) - -For larger models or multi-GPU environments, customize the deploy YAML (see [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)) and set per-stage `tensor_parallel_size` and `devices`: - -```yaml -# Example: TP=2 on GPUs 0,1 for the Thinker stage -stages: - - stage_id: 0 - tensor_parallel_size: 2 - devices: "0,1" -``` - -### VAE Patch Parallelism - -[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) splits Bagel VAE **decode/encode** tiles across multiple GPUs on the **DiT stage**, reducing **per-GPU peak memory during VAE decode**. Use it when high-resolution `text2img` or `img2img` hits VAE OOM or large decode spikes. - -**Bagel-specific notes:** - -- Implemented in `BagelPipeline` via `DistributedAutoEncoder` (DiT stage only). -- **Single-stage** is the simplest path: one DiT process with TP + VAE patch parallel. -- **Two-stage**: enable on **stage 1 (DiT)** only; stage 0 (Thinker) keeps encoder-only `VAEEncoder` and does not use VAE patch parallel. -- You need a DiT `world_size` ≥ `vae_patch_parallel_size` (typically `tensor_parallel_size=2` on that stage). VAE PP reuses the DiT process group; it is not a standalone second-GPU VAE worker. - -**Single-stage via deploy YAML** (recommended for `end2end.py`): - -```yaml -pipeline: bagel_single_stage -async_chunk: false - -stages: - - stage_id: 0 - max_num_batched_tokens: 32768 - max_num_seqs: 1 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false - devices: "0,1" - vae_use_tiling: true - parallel_config: - tensor_parallel_size: 2 - vae_patch_parallel_size: 2 - default_sampling_params: - seed: 52 -``` - -```bash -cd examples/offline_inference/bagel - -CUDA_VISIBLE_DEVICES=0,1 python end2end.py \ - --model /path/to/BAGEL-7B-MoT \ - --deploy-config /path/to/bagel_single_stage_vae_pp.yaml \ - --modality text2img \ - --prompts "A cute cat" \ - --steps 10 \ - --output ./out_vae_pp -``` - -**Single-stage via `Omni` kwargs** (same flags as online serving): - -```python -from vllm_omni.entrypoints.omni import Omni - -omni = Omni( - model="ByteDance-Seed/BAGEL-7B-MoT", - deploy_config="vllm_omni/deploy/bagel_single_stage.yaml", - tensor_parallel_size=2, - vae_patch_parallel_size=2, - vae_use_tiling=True, -) -# Then call omni.generate(...) as in end2end.py -``` - -**Two-stage (VAE PP on DiT only):** - -```yaml -stages: - - stage_id: 0 - devices: "0" - # AR Thinker — no vae_patch_parallel here - - - stage_id: 1 - devices: "0,1" - vae_use_tiling: true - parallel_config: - tensor_parallel_size: 2 - vae_patch_parallel_size: 2 -``` - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --deploy-config /path/to/bagel_vae_pp.yaml \ - --modality text2img \ - --prompts "A cute cat" -``` - -**Startup log checks:** - -```text -INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it. -``` - -| Setting | Role | -| :------ | :--- | -| `parallel_config.tensor_parallel_size` | DiT world size / TP (must be ≥ `vae_patch_parallel_size`) | -| `parallel_config.vae_patch_parallel_size` | Number of ranks for distributed VAE tiles (`1` = off) | -| `vae_use_tiling` | Enable spatial tiling (auto-enabled when `vae_patch_parallel_size > 1`) | - -#### Hybrid Sharded Data Parallel (HSDP) - -For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage. - -1. **Enable HSDP**: Set `use_hsdp: true`. -2. **Set shard size**: Set `hsdp_shard_size` to the number of GPUs used for sharding (for example, `4`). -3. **Set replicate size**: Usually keep `hsdp_replicate_size: 1` unless you want replicated HSDP groups. -4. **Set devices**: Specify the comma-separated GPU IDs used by the diffusion stage (for example, `"0,1,2,3"`). - -Example configuration for HSDP across 4 GPUs: -```yaml - - stage_id: 1 - devices: "0,1,2,3" - parallel_config: - use_hsdp: true - hsdp_shard_size: 4 - hsdp_replicate_size: 1 -``` - -Then pass the custom deploy YAML: - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A cute cat" \ - --deploy-config /path/to/custom_bagel.yaml -``` - -### FP8 Quantization - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A cute cat" \ - --quantization fp8 -``` - -## Command Line Reference - -### Core Arguments - -| Argument | Type | Default | Description | -| :------- | :--- | :------ | :---------- | -| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or HuggingFace name | -| `--modality` | choice | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` | -| `--prompts` | list | `None` | Input text prompts | -| `--txt-prompts` | string | `None` | Path to text file with one prompt per line | -| `--image-path` | string | `None` | Input image path (required for `img2img`/`img2text`) | -| `--output` | string | `.` | Output directory for saved images | -| `--steps` | int | `50` | Number of diffusion inference steps | -| `--seed` | int | `None` | Random seed for reproducibility | - -### Think Mode Arguments - -| Argument | Type | Default | Description | -| :------- | :--- | :------ | :---------- | -| `--think` | flag | `False` | Enable `...` planning/reasoning | -| `--max-think-tokens` | int | `1000` | Maximum tokens for think generation | -| `--do-sample` | flag | `False` | Use sampling instead of greedy decoding | -| `--text-temperature` | float | `0.3` | Sampling temperature for text generation | - -### CFG Arguments - -| Argument | Type | Default | Description | -| :------- | :--- | :------ | :---------- | -| `--cfg-text-scale` | float | `4.0` | Text CFG guidance scale | -| `--cfg-img-scale` | float | `1.5` | Image CFG guidance scale | -| `--negative-prompt` | string | `None` | Negative prompt for CFG | -| `--cfg-parallel-size` | int | `1` | CFG parallel GPU count (1, 2, or 3) | -| `--cfg-interval` | float[2] | pipeline default | CFG active window `[start, end]` | -| `--cfg-renorm-type` | string | `None` | `global`, `text_channel`, or `channel` | -| `--cfg-renorm-min` | float | `None` | Minimum renormalization value | - -### Engine Arguments - -| Argument | Type | Default | Description | -| :------- | :--- | :------ | :---------- | -| `--deploy-config` | string | `None` | Path to deploy YAML (auto-detected if omitted) | -| `--worker-backend` | choice | `process` | `process` or `ray` | -| `--ray-address` | string | `None` | Ray cluster address | -| `--quantization` | string | `None` | Quantization method (e.g. `fp8`) | -| `--log-stats` | flag | `False` | Enable statistics logging | -| `--init-timeout` | int | `300` | Initialization timeout (seconds) | -| `--batch-timeout` | int | `5` | Batch timeout (seconds) | -| `--enable-diffusion-pipeline-profiler` | flag | `False` | Profile diffusion stage durations | - -## FAQ - -- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML. - -**Two-stage VRAM usage:** - -| Stage | VRAM | -| :---- | :--- | -| Stage 0 (Thinker) | **15.04 GiB + KV Cache** | -| Stage 1 (DiT) | **26.50 GiB** | -| Total | **~42 GiB + KV Cache** | - -**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process. diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py deleted file mode 100644 index 088687253..000000000 --- a/examples/offline_inference/bagel/end2end.py +++ /dev/null @@ -1,278 +0,0 @@ -import argparse -import os - -from vllm_omni.inputs.data import OmniPromptType -from vllm_omni.model_executor.stage_input_processors.bagel import ( - GEN_THINK_SYSTEM_PROMPT, - VLM_THINK_SYSTEM_PROMPT, -) - - -def parse_args(): - parser = argparse.ArgumentParser() - parser.add_argument( - "--model", - default="ByteDance-Seed/BAGEL-7B-MoT", - help="Path to merged model directory.", - ) - parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.") - parser.add_argument( - "--txt-prompts", - type=str, - default=None, - help="Path to a .txt file with one prompt per line (preferred).", - ) - parser.add_argument("--prompt-type", default="text", choices=["text"]) - - parser.add_argument( - "--modality", - default="text2img", - choices=["text2img", "img2img", "img2text", "text2text"], - help="Modality mode to control stage execution.", - ) - - parser.add_argument( - "--image-path", - type=str, - default=None, - help="Path to input image for img2img.", - ) - - parser.add_argument( - "--output", - type=str, - default=".", - help="Output directory to save images.", - ) - - # OmniLLM init args - parser.add_argument("--log-stats", action="store_true", default=False) - parser.add_argument("--init-sleep-seconds", type=int, default=20) - parser.add_argument("--batch-timeout", type=int, default=5) - parser.add_argument("--init-timeout", type=int, default=300) - parser.add_argument("--shm-threshold-bytes", type=int, default=65536) - parser.add_argument("--worker-backend", type=str, default="process", choices=["process", "ray"]) - parser.add_argument("--ray-address", type=str, default=None) - parser.add_argument( - "--deploy-config", - type=str, - default=None, - help="Path to deploy YAML. If unset, auto-loads vllm_omni/deploy/bagel.yaml based on the HF model_type.", - ) - parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.") - - parser.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (default: 4.0)") - parser.add_argument("--cfg-img-scale", type=float, default=1.5, help="Image CFG scale (default: 1.5)") - parser.add_argument( - "--negative-prompt", type=str, default=None, help="Negative prompt for CFG (default: empty prompt)" - ) - parser.add_argument( - "--cfg-parallel-size", - type=int, - default=1, - choices=[1, 2, 3], - help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).", - ) - parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.") - parser.add_argument( - "--cfg-interval", - type=float, - nargs=2, - default=None, - help="CFG interval [start, end] (default: pipeline default)", - ) - parser.add_argument( - "--cfg-renorm-type", type=str, default=None, help="CFG renorm type: global, text_channel, channel" - ) - parser.add_argument("--cfg-renorm-min", type=float, default=None, help="CFG renorm min") - parser.add_argument( - "--enable-diffusion-pipeline-profiler", - action="store_true", - help="Enable diffusion pipeline profiler to display stage durations.", - ) - parser.add_argument( - "--quantization", - type=str, - default=None, - help="Quantization method (e.g. 'fp8').", - ) - parser.add_argument( - "--think", - action="store_true", - default=False, - help="Enable thinking mode: AR stage decodes ... planning tokens before image generation.", - ) - parser.add_argument( - "--max-think-tokens", - type=int, - default=1000, - help="Maximum number of tokens for thinking text generation (default: 1000).", - ) - parser.add_argument( - "--do-sample", - action="store_true", - default=False, - help="Enable sampling for text generation (default: greedy).", - ) - parser.add_argument( - "--text-temperature", - type=float, - default=0.3, - help="Temperature for text generation sampling (default: 0.3).", - ) - - args = parser.parse_args() - return args - - -def main(): - args = parse_args() - os.makedirs(args.output, exist_ok=True) - model_name = args.model - prompts: list[OmniPromptType] = [] - try: - if getattr(args, "txt_prompts", None) and args.prompt_type == "text": - with open(args.txt_prompts, encoding="utf-8") as f: - lines = [ln.strip() for ln in f.readlines()] - prompts = [ln for ln in lines if ln != ""] - print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}") - else: - prompts = args.prompts - except Exception as e: - print(f"[Error] Failed to load prompts: {e}") - raise - - if not prompts: - prompts = ["A cute cat"] - print(f"[Info] No prompts provided, using default: {prompts}") - - from PIL import Image - - from vllm_omni.entrypoints.omni import Omni - - omni_kwargs = vars(args).copy() - omni_kwargs.pop("model", None) - deploy_config = args.deploy_config - if args.think and deploy_config is None: - deploy_config = "vllm_omni/deploy/bagel_think.yaml" - print(f"[Info] Think mode enabled, using deploy config: {deploy_config}") - if deploy_config: - omni_kwargs["deploy_config"] = deploy_config - - if args.quantization: - omni_kwargs["quantization_config"] = args.quantization - - omni = Omni(model=model_name, **omni_kwargs) - - formatted_prompts = [] - for p in prompts: - if args.modality == "img2img": - if not args.image_path or not os.path.exists(args.image_path): - raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {args.image_path}") - loaded_image = Image.open(args.image_path).convert("RGB") - think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else "" - final_prompt_text = f"{think_prefix}<|fim_middle|><|im_start|>{p}<|im_end|>" - prompt_dict = { - "prompt": final_prompt_text, - "multi_modal_data": {"img2img": loaded_image}, - "modalities": ["img2img"], - } - if args.negative_prompt is not None: - prompt_dict["negative_prompt"] = args.negative_prompt - formatted_prompts.append(prompt_dict) - elif args.modality == "img2text": - if args.image_path: - loaded_image = Image.open(args.image_path).convert("RGB") - think_prefix = f"<|im_start|>system\n{VLM_THINK_SYSTEM_PROMPT}<|im_end|>\n" if args.think else "" - final_prompt_text = ( - f"{think_prefix}<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n" - ) - prompt_dict = { - "prompt": final_prompt_text, - "multi_modal_data": {"image": loaded_image}, - "modalities": ["text"], - } - formatted_prompts.append(prompt_dict) - elif args.modality == "text2text": - think_prefix = f"<|im_start|>{VLM_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else "" - final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|><|im_start|>" - prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]} - formatted_prompts.append(prompt_dict) - else: - think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else "" - final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|>" - prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]} - if args.negative_prompt is not None: - prompt_dict["negative_prompt"] = args.negative_prompt - formatted_prompts.append(prompt_dict) - - params_list = omni.default_sampling_params_list - # Bagel exposes 1 sampling param set for single-stage (DiT-only) and - # 2 for two-stage (Thinker + DiT). This heuristic may need updating - # if future pipelines break that 1:1 mapping. - is_single_stage = len(params_list) == 1 - - diffusion_params_idx = 0 if is_single_stage else (1 if len(params_list) > 1 else 0) - diffusion_params = params_list[diffusion_params_idx] - - if args.modality in ("text2img", "img2img"): - diffusion_params.num_inference_steps = args.steps # type: ignore - diffusion_params.cfg_parallel_size = args.cfg_parallel_size # type: ignore - if args.seed is not None: - diffusion_params.seed = args.seed # type: ignore - - extra = getattr(diffusion_params, "extra_args", {}) or {} - extra["cfg_text_scale"] = args.cfg_text_scale - extra["cfg_img_scale"] = args.cfg_img_scale - if args.cfg_interval is not None: - extra["cfg_interval"] = tuple(args.cfg_interval) - if args.cfg_renorm_type is not None: - extra["cfg_renorm_type"] = args.cfg_renorm_type - if args.cfg_renorm_min is not None: - extra["cfg_renorm_min"] = args.cfg_renorm_min - if args.negative_prompt is not None: - extra["negative_prompt"] = args.negative_prompt - - needs_text_gen = is_single_stage and (args.think or args.modality in ("text2text", "img2text")) - if needs_text_gen: - if args.think: - extra["think"] = True - extra["max_think_tokens"] = args.max_think_tokens - extra["do_sample"] = args.do_sample - extra["text_temperature"] = args.text_temperature - diffusion_params.extra_args = extra # type: ignore - - omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list)) - - img_idx = 0 - for req_output in omni_outputs: - # 2-stage think mode: text output from thinker stage - ro = getattr(req_output, "request_output", None) - if ro and getattr(ro, "outputs", None): - txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs) - if txt: - if args.think: - print(f"[Think]\n{txt}") - else: - print(f"[Output] Text:\n{txt}") - - # Single-stage DiT: text from custom_output - custom = getattr(req_output, "_custom_output", {}) or {} - if custom.get("think_text"): - print(f"[Think]\n{custom['think_text']}") - if custom.get("text_output"): - print(f"[Output] Text:\n{custom['text_output']}") - - images = getattr(req_output, "images", None) - if not images: - continue - - for j, img in enumerate(images): - save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png") - img.save(save_path) - print(f"[Output] Saved image to {save_path}") - img_idx += 1 - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 0f2445673..299efb33b 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -97,9 +97,16 @@ import torch from PIL import Image from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni +from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params from vllm_omni.inputs.data import OmniDiffusionSamplingParams -from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.model_extras import ( + build_image_to_image_prompt, + get_extra_body_params, + get_model_class_name, + should_init_extra_args_for_non_diffusion_stages, +) from vllm_omni.platforms import current_omni_platform @@ -187,6 +194,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--guidance-scale-2", type=float, default=None, help="image guidance scale for image-to-image generation." ) + parser.add_argument( + "--extra-args", + type=parse_profiler_config, + default=None, + help="JSON object copied to OmniDiffusionSamplingParams.extra_args, e.g. '{\"cfg_text_scale\": 4.0}'.", + ) parser.add_argument( "--output", type=str, @@ -468,6 +481,8 @@ def main(): enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, profiler_config=args.profiler_config, ) + model_class_name = get_model_class_name(omni) + declared_extra_body_params = get_extra_body_params(model_class_name) print("Pipeline loaded") profiler_enabled = args.profiler_config is not None @@ -478,6 +493,8 @@ def main(): print(f" Model: {args.model}") print(f" Inference steps: {args.num_inference_steps}") print(f" Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}") + if args.height is not None or args.width is not None: + print(f" Output size: {args.width or 'auto'}x{args.height or 'auto'}") if isinstance(input_image, list): print(f" Number of input images: {len(input_image)}") for idx, img in enumerate(input_image): @@ -487,7 +504,6 @@ def main(): 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}, enable_expert_parallel: {args.enable_expert_parallel}" ) - print(f" Output image size: {args.width or 'default'}x{args.height or 'default'}") print(f"{'=' * 60}\n") generation_start = time.perf_counter() @@ -496,10 +512,20 @@ def main(): print("[Profiler] Starting profiling...") omni.start_profile() - # Generate edited image - sampling_params = OmniDiffusionSamplingParams( - width=args.width, + prompt_dict = build_image_to_image_prompt( + model_class_name=model_class_name, + prompt=args.prompt, + negative_prompt=args.negative_prompt, + input_image=input_image, height=args.height, + width=args.width, + ) + + extra_args_from_cli = dict(args.extra_args or {}) + if args.negative_prompt is not None: + extra_args_from_cli.setdefault("negative_prompt", args.negative_prompt) + + diffusion_params = OmniDiffusionSamplingParams( generator=generator, true_cfg_scale=args.cfg_scale, guidance_scale=args.guidance_scale, @@ -507,17 +533,44 @@ def main(): num_inference_steps=args.num_inference_steps, num_outputs_per_prompt=args.num_outputs_per_prompt, layers=args.layers, + resolution=args.resolution, + height=args.height, + width=args.width, ) - if args.resolution: - sampling_params.resolution = args.resolution - outputs = omni.generate( - { - "prompt": args.prompt, - "negative_prompt": args.negative_prompt, - "multi_modal_data": {"image": input_image}, - }, - sampling_params, + if declared_extra_body_params: + apply_declared_extra_args( + diffusion_params, + declared_extra_body_params, + extra_args_from_cli, + ) + else: + diffusion_params.extra_args.update({k: v for k, v in extra_args_from_cli.items() if v is not None}) + + # Build per-stage sampling params for multi-stage models + init_non_diffusion = should_init_extra_args_for_non_diffusion_stages( + model_class_name, ) + defaults = list(omni.default_sampling_params_list or []) + sampling_params_list = [clone_sampling_params(p) for p in defaults] + if not sampling_params_list: + sampling_params_list = [diffusion_params] + + diffusion_replaced = False + for idx, params in enumerate(sampling_params_list): + if isinstance(params, OmniDiffusionSamplingParams): + merged_extra = dict(getattr(params, "extra_args", {}) or {}) + merged_extra.update(diffusion_params.extra_args) + diffusion_params.extra_args = merged_extra + sampling_params_list[idx] = diffusion_params + diffusion_replaced = True + elif init_non_diffusion and hasattr(params, "extra_args"): + if params.extra_args is None: + params.extra_args = {} + + if not diffusion_replaced and len(sampling_params_list) == 1: + sampling_params_list = [diffusion_params] + + outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list) generation_end = time.perf_counter() generation_time = generation_end - generation_start @@ -544,17 +597,16 @@ def main(): if not outputs: raise ValueError("No output generated from omni.generate()") - # Extract images from OmniRequestOutput - # omni.generate() returns list[OmniRequestOutput], extract images from request_output.images - first_output = outputs[0] - if not hasattr(first_output, "request_output") or not first_output.request_output: - raise ValueError("No request_output found in OmniRequestOutput") + images = None + for output in outputs: + images = getattr(output, "images", None) + if images: + break + req_out = getattr(output, "request_output", None) + images = getattr(req_out, "images", None) if req_out is not None else None + if images: + break - req_out = first_output.request_output - if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"): - raise ValueError("Invalid request_output structure or missing 'images' key") - - images = req_out.images if not images: raise ValueError("No images found in request_output") diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index b65ffa0da..06e5dc181 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -10,10 +10,18 @@ from typing import Any import torch from vllm_omni.diffusion.data import logger +from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni +from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.lora.request import LoRARequest from vllm_omni.lora.utils import stable_lora_int_id +from vllm_omni.model_extras import ( + build_text_to_image_prompt, + get_extra_body_params, + get_model_class_name, + should_init_extra_args_for_non_diffusion_stages, +) from vllm_omni.platforms import current_omni_platform @@ -259,6 +267,19 @@ def parse_args() -> argparse.Namespace: action="store_true", help="[NextStep-1.1 only] Apply layer normalization to sampled tokens.", ) + parser.add_argument( + "--extra-body", + type=parse_profiler_config, + default=None, + help=( + "Model-specific generation params as a JSON object, e.g. " + '\'{"timestep_shift": 3.0, "cfg_text_scale": 4.0, "cfg_interval": [0.4, 1.0]}\'. ' + "Each key is filtered against the model's declared extra_body_params " + "(see vllm_omni/model_extras), so unknown keys for the chosen model are " + "silently dropped. Values here take precedence over the equivalent " + "model-specific flags above." + ), + ) parser.add_argument( "--enable-diffusion-pipeline-profiler", action="store_true", @@ -405,6 +426,8 @@ def main(): # NextStep-1.1 requires explicit pipeline class omni_kwargs["model_class_name"] = "NextStep11Pipeline" omni = Omni(**omni_kwargs) + model_class_name = get_model_class_name(omni) + declared_extra_body_params = get_extra_body_params(model_class_name) if profiler_enabled: print("[Profiler] Starting profiling...") @@ -445,34 +468,79 @@ def main(): ) generation_start = time.perf_counter() - extra_args = { + + prompt_dict = build_text_to_image_prompt( + model_class_name=model_class_name, + prompt=args.prompt, + negative_prompt=args.negative_prompt, + height=args.height, + width=args.width, + ) + + diffusion_params = OmniDiffusionSamplingParams( + height=args.height, + width=args.width, + seed=args.seed, + generator=generator, + true_cfg_scale=args.cfg_scale, + guidance_scale=args.guidance_scale, + guidance_scale_2=args.guidance_scale_2, + num_inference_steps=args.num_inference_steps, + num_outputs_per_prompt=args.num_images_per_prompt, + ) + + # Base layer: backward-compatible model-specific flags. New model params should + # instead be declared in vllm_omni/model_extras and passed via --extra-body, so + # this dict does not need to grow per-model. + user_extra = { + "cfg_scale": args.cfg_scale, + "cfg_text_scale": args.cfg_scale, + "negative_prompt": args.negative_prompt, + "timestep_shift": args.timesteps_shift, "timesteps_shift": args.timesteps_shift, "cfg_schedule": args.cfg_schedule, "use_norm": args.use_norm, "use_system_prompt": args.use_system_prompt, "system_prompt": args.system_prompt, } - if lora_request: - extra_args["lora_request"] = lora_request - extra_args["lora_scale"] = args.lora_scale + # Override layer: generic JSON passthrough wins over the flags above. Keys are + # still filtered against the model's declared extra_body_params downstream. + if args.extra_body: + user_extra.update(args.extra_body) + if declared_extra_body_params: + apply_declared_extra_args(diffusion_params, declared_extra_body_params, user_extra) + else: + diffusion_params.extra_args.update({k: v for k, v in user_extra.items() if v is not None}) - outputs = omni.generate( - { - "prompt": args.prompt, - "negative_prompt": args.negative_prompt, - }, - OmniDiffusionSamplingParams( - height=args.height, - width=args.width, - generator=generator, - true_cfg_scale=args.cfg_scale, - guidance_scale=args.guidance_scale, - guidance_scale_2=args.guidance_scale_2, - num_inference_steps=args.num_inference_steps, - num_outputs_per_prompt=args.num_images_per_prompt, - extra_args=extra_args, - ), + if lora_request: + diffusion_params.extra_args["lora_request"] = lora_request + diffusion_params.extra_args["lora_scale"] = args.lora_scale + + # Build per-stage sampling params for multi-stage models (e.g. BAGEL), + # or wrap single diffusion params for single-stage models. + init_non_diffusion = should_init_extra_args_for_non_diffusion_stages( + model_class_name, ) + defaults = list(omni.default_sampling_params_list or []) + sampling_params_list = [clone_sampling_params(p) for p in defaults] + if not sampling_params_list: + sampling_params_list = [diffusion_params] + + diffusion_replaced = False + for idx, params in enumerate(sampling_params_list): + if isinstance(params, OmniDiffusionSamplingParams): + sampling_params_list[idx] = diffusion_params + diffusion_replaced = True + elif init_non_diffusion and hasattr(params, "extra_args"): + if params.extra_args is None: + params.extra_args = {} + if args.seed is not None and hasattr(params, "seed"): + params.seed = args.seed + + if not diffusion_replaced and len(sampling_params_list) == 1: + sampling_params_list = [diffusion_params] + + outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list) generation_end = time.perf_counter() generation_time = generation_end - generation_start @@ -502,15 +570,16 @@ def main(): raise ValueError("No output generated from omni.generate()") logger.info(f"Outputs: {outputs}") - first_output = outputs[0] - if not hasattr(first_output, "request_output") or not first_output.request_output: - raise ValueError("No request_output found in OmniRequestOutput") + images = None + for output in outputs: + images = getattr(output, "images", None) + if images: + break + req_out = getattr(output, "request_output", None) + images = getattr(req_out, "images", None) if req_out is not None else None + if images: + break - req_out = first_output.request_output - if not hasattr(req_out, "images"): - raise ValueError("Invalid request_output structure or missing 'images'.") - - images = req_out.images if not images: raise ValueError("No images found in request_output") diff --git a/examples/online_serving/bagel/README.md b/examples/online_serving/bagel/README.md deleted file mode 100644 index b7a54da35..000000000 --- a/examples/online_serving/bagel/README.md +++ /dev/null @@ -1,404 +0,0 @@ -# BAGEL-7B-MoT - -## Installation - -Please refer to [README.md](../../../README.md) - -## Architecture - -BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies: - -| Topology | Stages | Description | -| :------- | :----- | :---------- | -| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. | -| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. | - -Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`. - -> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices. - -## Launch the Server - -### Two-Stage (Default) - -The default pipeline is auto-detected from the model. No extra flags needed: - -```bash -vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 -``` - -Or use the convenience script: - -```bash -cd examples/online_serving/bagel -bash run_server.sh - -# Initialize each stage in a discrete isolated process terminal -bash run_server_stage_cli.sh --stage 0 -bash run_server_stage_cli.sh --stage 1 -``` - -To use a custom deploy YAML, pass it via `--deploy-config`: - -```bash -vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \ - --deploy-config /path/to/deploy_config.yaml -``` - -See [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml) for the default two-stage deploy configuration. - -### Single-Stage - -The DiT stage contains a full LLM, ViT, VAE, and tokenizer, so it can handle all modalities (text2img, img2img, img2text, text2text, think) without a separate Thinker stage: - -```bash -vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \ - --deploy-config vllm_omni/deploy/bagel_single_stage.yaml -``` - -See [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) for configuration. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry. - -### Tensor Parallelism (TP) - -For larger models or multi-GPU environments, enable TP via CLI: - -```bash -vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size 2 -``` - -Or set `tensor_parallel_size` per stage in a custom deploy YAML. - -### VAE Patch Parallelism - -[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) distributes Bagel VAE **decode/encode** across multiple GPUs by splitting latent tiles. It lowers **per-GPU peak memory during VAE decode**, which helps high-resolution `text2img` / `img2img` when VAE becomes a bottleneck. - -**Scope for Bagel:** - -| Topology | VAE patch parallel | -| :------- | :----------------- | -| **Single-stage** (DiT only) | Supported on stage 0 (`BagelPipeline` + `DistributedAutoEncoder`) | -| **Two-stage** | Supported on **stage 1 (DiT)** only; stage 0 (Thinker) uses encoder-only VAE and is unrelated | - -**Requirements:** - -- `vae_patch_parallel_size > 1` and a distributed VAE (`DistributedAutoEncoder` on the DiT pipeline). -- The DiT process group must have at least `vae_patch_parallel_size` ranks. In practice this means the diffusion stage `world_size` must be ≥ 2 (commonly `tensor_parallel_size=2` on that stage). -- `vae_use_tiling` must be enabled. If you set `vae_patch_parallel_size > 1` and omit tiling, the registry auto-enables `vae_use_tiling` at startup. - -VAE patch parallel **reuses the DiT process group** (`dit_group`); it does not create a separate VAE-only worker pool. It is not a substitute for single-GPU VAE tiling (`vae_pp=1`). - -**Online serving (single-stage, 2 GPUs):** - -```bash -CUDA_VISIBLE_DEVICES=0,1 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \ - --deploy-config vllm_omni/deploy/bagel_single_stage.yaml \ - --tensor-parallel-size 2 \ - --vae-patch-parallel-size 2 \ - --vae-use-tiling -``` - -**Online serving (two-stage, VAE PP on DiT stage 1):** use a custom deploy YAML, for example: - -```yaml -stages: - - stage_id: 0 - devices: "0" - # Thinker (AR) — no VAE patch parallel here - - - stage_id: 1 - devices: "0,1" - vae_use_tiling: true - parallel_config: - tensor_parallel_size: 2 - vae_patch_parallel_size: 2 -``` - -```bash -vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \ - --deploy-config /path/to/bagel_vae_pp.yaml -``` - -**Verify it is active** (check server logs at startup): - -```text -INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it. -``` - -| CLI flag | Default | Description | -| :------- | :------ | :---------- | -| `--vae-patch-parallel-size` | `1` | Number of DiT ranks used for VAE tile parallelism. Set to `2` or higher to enable. Should be ≤ DiT process group size (typically match `--tensor-parallel-size` on the diffusion stage). | -| `--vae-use-tiling` | off | Enable VAE spatial tiling. Required for VAE patch parallel (auto-enabled when `vae_patch_parallel_size > 1`). | - -#### Hybrid Sharded Data Parallel (HSDP) - -For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage. - -1. **Enable HSDP**: Set `use_hsdp: true`. -2. **Set shard size**: Set `hsdp_shard_size` to the number of GPUs used for sharding (for example, `4`). -3. **Set replicate size**: Usually keep `hsdp_replicate_size: 1` unless you want replicated HSDP groups. -4. **Set devices**: Specify the comma-separated GPU IDs used by the diffusion stage (for example, `"0,1,2,3"`). - -Example configuration for HSDP across 4 GPUs: -```yaml - - stage_id: 1 - devices: "0,1,2,3" - parallel_config: - use_hsdp: true - hsdp_shard_size: 4 - hsdp_replicate_size: 1 -``` - -### Multi-Node Deployment - -Deploy each stage on a **separate node** for better resource utilization. Replace `` with the actual IP address of your orchestrator node. - -**1. Launch Stage 0 (Thinker / Orchestrator)** on the orchestrator node: - -```bash -# API server port for client requests: 8000 -vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \ - --port 8000 \ - --stage-id 0 \ - --omni-master-address \ - --omni-master-port 8091 -``` - -**2. Launch Stage 1 (DiT)** on the remote node in headless mode: - -```bash -vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \ - --stage-id 1 \ - --headless \ - --omni-master-address \ - --omni-master-port 8091 -``` - -Or use the convenience script: - -```bash -# Terminal 1: Stage 0 -bash run_server_stage_cli.sh --stage 0 - -# Terminal 2: Stage 1 -bash run_server_stage_cli.sh --stage 1 - -# With extra args -bash run_server_stage_cli.sh --stage 0 -- --tensor-parallel-size 2 -bash run_server_stage_cli.sh --stage 1 -- --gpu-memory-utilization 0.9 -``` - -**vllm serve arguments:** - -| Argument | Description | -| :------- | :---------- | -| `--stage-id` | Which stage this process runs (0 = Thinker, 1 = DiT) | -| `--headless` | Run without the API server (worker-only mode) | -| `-oma` / `--omni-master-address` | Orchestrator master address | -| `-omp` / `--omni-master-port` | Orchestrator master port | - -> [!IMPORTANT] -> **Startup Order**: Stage 0 (orchestrator) must be launched **before** Stage 1 (headless). -> Stage 0 will appear to hang on startup until Stage 1 (worker) connects — this is expected behavior. - -### Inter-Stage Connectors - -When deploying stages across nodes, configure the connector type in the deploy YAML: - -- **SharedMemoryConnector** (default): Used for single-node deployments. No explicit configuration needed. -- **MooncakeTransferEngineConnector**: For multi-node setups with RDMA hardware. Defined in [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml) under `connectors.rdma_connector`. - -To use Mooncake, create a custom deploy YAML that binds `output_connectors` / `input_connectors` on each stage to the `rdma_connector` defined in the `connectors` section. - -## Send Requests - -```bash -cd examples/online_serving/bagel -``` - -### Text to Image (text2img) - -**Python client:** - -```bash -python openai_chat_client.py \ - --prompt "A beautiful sunset over mountains" \ - --modality text2img \ - --output sunset.png \ - --steps 50 -``` - -**curl:** - -```bash -curl http://localhost:8091/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>A beautiful sunset over mountains<|im_end|>"}]}], - "modalities": ["image"], - "height": 512, - "width": 512, - "num_inference_steps": 50, - "seed": 42 - }' -``` - -### Image to Image (img2img) - -**Python client:** - -```bash -python openai_chat_client.py \ - --prompt "Make the cat stand up" \ - --modality img2img \ - --image-url /path/to/input.jpg \ - --output transformed.png -``` - -**curl:** - -```bash -IMAGE_BASE64=$(base64 -w 0 cat.jpg) - -cat < payload.json -{ - "messages": [{ - "role": "user", - "content": [ - {"type": "text", "text": "<|im_start|>Make the cat stand up<|im_end|>"}, - {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}} - ] - }], - "modalities": ["image"], - "height": 512, - "width": 512, - "num_inference_steps": 50, - "seed": 42 -} -EOF - -curl http://localhost:8091/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d @payload.json -``` - -### Image to Text (img2text) - -**Python client:** - -```bash -python openai_chat_client.py \ - --prompt "Describe this image in detail" \ - --modality img2text \ - --image-url /path/to/image.jpg -``` - -**curl:** - -```bash -IMAGE_BASE64=$(base64 -w 0 cat.jpg) - -cat < payload.json -{ - "messages": [{ - "role": "user", - "content": [ - {"type": "text", "text": "<|im_start|>user\n<|image_pad|>\nDescribe this image in detail<|im_end|>\n<|im_start|>assistant\n"}, - {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}} - ] - }], - "modalities": ["text"] -} -EOF - -curl http://localhost:8091/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d @payload.json -``` - -### Text to Text (text2text) - -**Python client:** - -```bash -python openai_chat_client.py \ - --prompt "What is the capital of France?" \ - --modality text2text -``` - -**curl:** - -```bash -curl http://localhost:8091/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}], - "modalities": ["text"] - }' -``` - -### Python Client Arguments - -| Argument | Default | Description | -| :------- | :------ | :---------- | -| `--prompt` / `-p` | `A cute cat` | Text prompt | -| `--output` / `-o` | `bagel_output.png` | Output file path | -| `--server` / `-s` | `http://localhost:8091` | Server URL | -| `--image-url` / `-i` | `None` | Input image URL or local path (img2img/img2text) | -| `--modality` / `-m` | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` | -| `--height` | `512` | Image height in pixels | -| `--width` | `512` | Image width in pixels | -| `--steps` | `25` | Number of inference steps | -| `--seed` | `42` | Random seed | -| `--negative` | `None` | Negative prompt for CFG | - -Example with custom parameters: - -```bash -python openai_chat_client.py \ - --prompt "A futuristic city" \ - --modality text2img \ - --height 768 \ - --width 768 \ - --steps 50 \ - --seed 42 \ - --negative "blurry, low quality" -``` - -## Configuration Reference - -### Deploy YAML Files - -| File | Description | -| :--- | :---------- | -| [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml) | Two-stage default (Thinker + DiT on GPU 0) | -| [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) | Single-stage (DiT only) | - -### Key Deploy YAML Fields - -| Field | Scope | Description | -| :---- | :---- | :---------- | -| `pipeline` | top-level | Override auto-detected pipeline (e.g. `bagel_single_stage`) | -| `stages[].stage_id` | per-stage | Stage identifier (0, 1, ...) | -| `stages[].devices` | per-stage | GPU device IDs (e.g. `"0"`, `"0,1"`) | -| `stages[].max_num_seqs` | per-stage | Maximum concurrent sequences | -| `stages[].gpu_memory_utilization` | per-stage | Fraction of GPU memory to use | -| `stages[].enforce_eager` | per-stage | Disable CUDA graphs | -| `stages[].tensor_parallel_size` | per-stage | TP degree for this stage | -| `stages[].parallel_config.vae_patch_parallel_size` | per-stage (DiT) | VAE tile parallelism degree (DiT stage only) | -| `stages[].vae_use_tiling` | per-stage (DiT) | Enable VAE tiling (required for VAE patch parallel) | -| `connectors` | top-level | Define available connector instances (SHM, Mooncake) | -| `platforms` | top-level | Platform-specific overrides (e.g. `xpu`) | - -## FAQ - -- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML. - -**Two-stage VRAM usage:** - -| Stage | VRAM | -| :---- | :--- | -| Stage 0 (Thinker) | **15.04 GiB + KV Cache** | -| Stage 1 (DiT) | **26.50 GiB** | -| Total | **~42 GiB + KV Cache** | - -**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process. diff --git a/examples/online_serving/bagel/openai_chat_client.py b/examples/online_serving/bagel/openai_chat_client.py deleted file mode 100755 index cc9ec32db..000000000 --- a/examples/online_serving/bagel/openai_chat_client.py +++ /dev/null @@ -1,184 +0,0 @@ -#!/usr/bin/env python3 -""" -Bagel OpenAI-compatible chat client for image generation and multimodal tasks. - -Usage: - python openai_chat_client.py --prompt "A cute cat" --output output.png - python openai_chat_client.py --prompt "Describe this image" --image-url https://example.com/image.png -""" - -import argparse -import base64 -from pathlib import Path - -import requests - - -def generate_image( - prompt: str, - server_url: str = "http://localhost:8091", - image_url: str | None = None, - height: int | None = None, - width: int | None = None, - steps: int | None = None, - seed: int | None = None, - negative_prompt: str | None = None, - modality: str = "text2img", # "text2img" (default), "img2img", "img2text", "text2text" -) -> bytes | str | None: - """Generate an image or text using the chat completions API. - - Args: - prompt: Text description or prompt - server_url: Server URL - image_url: URL or path to input image (for img2img/img2text) - height: Image height in pixels - width: Image width in pixels - steps: Number of inference steps - seed: Random seed - negative_prompt: Negative prompt - modality: Task modality hint - - Returns: - Image bytes (for image outputs) or Text string (for text outputs) or None if failed - """ - - # Construct Message Content - content = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}] - - if image_url: - # Check if local file - if Path(image_url).exists(): - with open(image_url, "rb") as f: - b64_data = base64.b64encode(f.read()).decode("utf-8") - final_image_url = f"data:image/jpeg;base64,{b64_data}" - else: - final_image_url = image_url - - content.append({"type": "image_url", "image_url": {"url": final_image_url}}) - - messages = [{"role": "user", "content": content}] - - # Build request payload with all parameters at top level - # Note: vLLM ignores "extra_body", so we put parameters directly in the payload - payload = {"messages": messages} - - # Set output modalities at top level - if modality == "text2img" or modality == "img2img": - payload["modalities"] = ["image"] - elif modality == "img2text" or modality == "text2text": - payload["modalities"] = ["text"] - - # Add generation parameters directly to payload - if height is not None: - payload["height"] = height - if width is not None: - payload["width"] = width - if steps is not None: - payload["num_inference_steps"] = steps - if seed is not None: - payload["seed"] = seed - if negative_prompt: - payload["negative_prompt"] = negative_prompt - - # Send request - try: - print(f"Sending request to {server_url} with modality {modality}...") - response = requests.post( - f"{server_url}/v1/chat/completions", - headers={"Content-Type": "application/json"}, - json=payload, - timeout=300, - ) - response.raise_for_status() - data = response.json() - - # Extract content - check ALL choices since server may return multiple - # (e.g., text in choices[0], image in choices[1]) - choices = data.get("choices", []) - - # First pass: look for image output in any choice - for choice in choices: - choice_content = choice.get("message", {}).get("content") - - # Handle Image Output - if isinstance(choice_content, list) and len(choice_content) > 0: - first_item = choice_content[0] - if isinstance(first_item, dict) and "image_url" in first_item: - img_url_str = first_item["image_url"].get("url", "") - if img_url_str.startswith("data:image"): - _, b64_data = img_url_str.split(",", 1) - return base64.b64decode(b64_data) - - # Second pass: look for text output if no image found - for choice in choices: - choice_content = choice.get("message", {}).get("content") - if isinstance(choice_content, str) and choice_content: - return choice_content - - print(f"Unexpected response format: {choices}") - return None - - except Exception as e: - print(f"Error: {e}") - return None - - -def main(): - parser = argparse.ArgumentParser(description="Bagel multimodal chat client") - parser.add_argument("--prompt", "-p", default="A cute cat", help="Text prompt") - parser.add_argument("--output", "-o", default="bagel_output.png", help="Output file (for image results)") - parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL") - - # Modality Control - parser.add_argument("--image-url", "-i", type=str, help="Input image URL or local path") - parser.add_argument( - "--modality", - "-m", - default="text2img", - choices=["text2img", "img2img", "img2text", "text2text"], - help="Task modality", - ) - - # Generation Params - parser.add_argument("--height", type=int, default=512, help="Image height") - parser.add_argument("--width", type=int, default=512, help="Image width") - parser.add_argument("--steps", type=int, default=25, help="Inference steps") - parser.add_argument("--seed", type=int, default=42, help="Random seed") - parser.add_argument("--negative", help="Negative prompt") - - args = parser.parse_args() - - print(f"Mode: {args.modality}") - if args.image_url: - print(f"Input Image: {args.image_url}") - - result = generate_image( - prompt=args.prompt, - server_url=args.server, - image_url=args.image_url, - height=args.height, - width=args.width, - steps=args.steps, - seed=args.seed, - negative_prompt=args.negative, - modality=args.modality, - ) - - if result: - if isinstance(result, bytes): - # It's an image - output_path = Path(args.output) - output_path.write_bytes(result) - print(f"Image saved to: {output_path}") - print(f"Size: {len(result) / 1024:.1f} KB") - elif isinstance(result, str): - # It's text - print("Response:") - print(result) - else: - print("Failed to generate response") - exit(1) - - -if __name__ == "__main__": - main() diff --git a/examples/online_serving/bagel/run_server.sh b/examples/online_serving/bagel/run_server.sh deleted file mode 100755 index a64057ef0..000000000 --- a/examples/online_serving/bagel/run_server.sh +++ /dev/null @@ -1,12 +0,0 @@ -#!/bin/bash -# Bagel online serving startup script - -MODEL="${MODEL:-ByteDance-Seed/BAGEL-7B-MoT}" -PORT="${PORT:-8091}" - -echo "Starting Bagel server..." -echo "Model: $MODEL" -echo "Port: $PORT" - -vllm serve "$MODEL" --omni \ - --port "$PORT" diff --git a/examples/online_serving/bagel/run_server_stage_cli.sh b/examples/online_serving/bagel/run_server_stage_cli.sh deleted file mode 100644 index 912e212f9..000000000 --- a/examples/online_serving/bagel/run_server_stage_cli.sh +++ /dev/null @@ -1,164 +0,0 @@ -#!/bin/bash -# Bagel multi-stage online serving startup script. -# -# Usage: -# ./run_server_stage_cli.sh --stage 0 -# ./run_server_stage_cli.sh --stage 1 -# ./run_server_stage_cli.sh --stage 0 -- --tensor-parallel-size 2 -# ./run_server_stage_cli.sh --stage 1 -- --gpu-memory-utilization 0.9 -# -# By default, `--stage all` keeps the old behavior and launches both stages in -# one session. Use `--stage 0` / `--stage 1` to launch each stage separately in -# different terminal sessions, with stage-specific extra CLI arguments passed -# after `--`. - -set -euo pipefail - -MODEL="${MODEL:-ByteDance-Seed/BAGEL-7B-MoT}" -PORT="${PORT:-8091}" -MASTER_ADDRESS="${MASTER_ADDRESS:-127.0.0.1}" -MASTER_PORT="${MASTER_PORT:-8092}" -STAGE="all" -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -DEPLOY_CONFIG="${DEPLOY_CONFIG:-$SCRIPT_DIR/../../../vllm_omni/deploy/bagel.yaml}" -EXTRA_ARGS=() - -usage() { - cat <&2 - usage - exit 1 - ;; - esac -done - -if [[ "$STAGE" != "0" && "$STAGE" != "1" && "$STAGE" != "all" ]]; then - echo "Invalid --stage value: $STAGE" >&2 - usage - exit 1 -fi - -print_config() { - echo "Model: $MODEL" - echo "API Port: $PORT" - echo "Master Address: $MASTER_ADDRESS" - echo "Master Port: $MASTER_PORT" - echo "Deploy Config: $DEPLOY_CONFIG" - echo "Selected Stage: $STAGE" - if [[ ${#EXTRA_ARGS[@]} -gt 0 ]]; then - echo "Extra Args: ${EXTRA_ARGS[*]}" - fi -} - -run_stage_0() { - echo "Starting Stage 0 (Thinker) as master..." - vllm serve "$MODEL" --omni \ - --port "$PORT" \ - --deploy-config "$DEPLOY_CONFIG" \ - --stage-id 0 \ - --omni-master-address "$MASTER_ADDRESS" \ - --omni-master-port "$MASTER_PORT" \ - "${EXTRA_ARGS[@]}" -} - -run_stage_1() { - echo "Starting Stage 1 (DiT) in headless mode..." - vllm serve "$MODEL" --omni \ - --deploy-config "$DEPLOY_CONFIG" \ - --stage-id 1 \ - --headless \ - --omni-master-address "$MASTER_ADDRESS" \ - --omni-master-port "$MASTER_PORT" \ - "${EXTRA_ARGS[@]}" -} - -echo "Starting Bagel multi-stage server..." -print_config - -case "$STAGE" in - 0) - run_stage_0 - ;; - 1) - run_stage_1 - ;; - all) - echo "Launching both stages in one session (legacy mode)..." - echo "Starting Stage 0 (Thinker) in background first..." - run_stage_0 & - STAGE_0_PID=$! - - cleanup() { - if [[ -n "${STAGE_0_PID:-}" ]]; then - kill "$STAGE_0_PID" 2>/dev/null || true - wait "$STAGE_0_PID" 2>/dev/null || true - fi - } - - trap cleanup EXIT INT TERM - - echo "Waiting briefly for Stage 0 to initialize..." - sleep 2 - run_stage_1 - ;; -esac diff --git a/examples/online_serving/image_to_image/README.md b/examples/online_serving/image_to_image/README.md index 59b1f0e2c..8a6c6a0e6 100644 --- a/examples/online_serving/image_to_image/README.md +++ b/examples/online_serving/image_to_image/README.md @@ -1,6 +1,8 @@ # Image-To-Image -This example demonstrates how to deploy Qwen-Image-Edit model for online image editing service using vLLM-Omni. +This example demonstrates how to deploy image-to-image models for online image editing service using vLLM-Omni. + +Supported models include Qwen-Image-Edit, BAGEL, and other image-to-image pipelines. For **multi-image** input editing, use **Qwen-Image-Edit-2509** (QwenImageEditPlusPipeline) and send multiple images in the user message content. @@ -21,6 +23,12 @@ vllm serve Qwen/Qwen-Image-Edit --omni --port 8092 vllm serve Qwen/Qwen-Image-Edit-2509 --omni --port 8092 ``` +### BAGEL + +```bash +vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 +``` + ### Start with Parameters @@ -119,6 +127,16 @@ python openai_chat_client.py --input input.png --prompt "Convert to oil painting python openai_chat_client.py --input input1.png input2.png --prompt "Combine these images into a single scene" --output output.png ``` +Pass model-specific parameters through `--extra-body` (e.g. for BAGEL): + +```bash +python openai_chat_client.py \ + --input input.png \ + --prompt "Make the scene look like a watercolor painting" \ + --server http://localhost:8091 \ + --extra-body '{"cfg_text_scale": 4.0, "cfg_img_scale": 1.5}' +``` + ### Method 4: Using Gradio Demo ```bash @@ -318,6 +336,9 @@ count, use `size` and `n` rather than `height`, `width`, or | `layers` | int | 4 | Number of layers (Qwen-Image-Layered) | | `resolution` | int | 640 | Resolution, 640 or 1024 (Qwen-Image-Layered) | +Models like BAGEL accept additional parameters via `extra_body` (e.g. +`cfg_text_scale`, `cfg_img_scale`). See the [BAGEL recipe](../../../recipes/Bagel/BAGEL-7B-MoT.md) for the full list. + ## Response Format ```json diff --git a/examples/online_serving/image_to_image/openai_chat_client.py b/examples/online_serving/image_to_image/openai_chat_client.py index 7aac2b214..46f2ad5a7 100644 --- a/examples/online_serving/image_to_image/openai_chat_client.py +++ b/examples/online_serving/image_to_image/openai_chat_client.py @@ -38,6 +38,7 @@ def edit_image( guidance_scale: float | None = None, seed: int | None = None, negative_prompt: str | None = None, + extra_body: dict | None = None, ) -> bytes | None: """Edit an image using the chat completions API. @@ -51,6 +52,7 @@ def edit_image( guidance_scale: CFG guidance scale seed: Random seed negative_prompt: Negative prompt + extra_body: Additional model-specific params (e.g. cfg_text_scale for BAGEL) Returns: Edited image bytes or None if failed @@ -75,24 +77,26 @@ def edit_image( ] # Build extra_body with generation parameters - extra_body = {} + merged_extra_body: dict[str, object] = {} if height is not None: - extra_body["height"] = height + merged_extra_body["height"] = height if width is not None: - extra_body["width"] = width + merged_extra_body["width"] = width if steps is not None: - extra_body["num_inference_steps"] = steps + merged_extra_body["num_inference_steps"] = steps if guidance_scale is not None: - extra_body["guidance_scale"] = guidance_scale + merged_extra_body["guidance_scale"] = guidance_scale if seed is not None: - extra_body["seed"] = seed + merged_extra_body["seed"] = seed if negative_prompt: - extra_body["negative_prompt"] = negative_prompt + merged_extra_body["negative_prompt"] = negative_prompt + if extra_body: + merged_extra_body.update(extra_body) # Build request payload - payload = {"messages": messages} - if extra_body: - payload["extra_body"] = extra_body + payload: dict[str, object] = {"messages": messages} + if merged_extra_body: + payload["extra_body"] = merged_extra_body # Send request try: @@ -121,6 +125,19 @@ def edit_image( return None +def parse_extra_body(value: str) -> dict: + """Parse a JSON string into a dict for --extra-body.""" + import json + + try: + obj = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"--extra-body must be valid JSON: {e}") from e + if not isinstance(obj, dict): + raise argparse.ArgumentTypeError("--extra-body must be a JSON object") + return obj + + def main(): parser = argparse.ArgumentParser(description="Qwen-Image-Edit chat client") parser.add_argument("--input", "-i", required=True, nargs="+", help="Input image path(s)") @@ -133,6 +150,12 @@ def main(): parser.add_argument("--guidance", type=float, default=7.5, help="Guidance scale") parser.add_argument("--seed", type=int, default=0, help="Random seed") parser.add_argument("--negative", help="Negative prompt") + parser.add_argument( + "--extra-body", + type=parse_extra_body, + default=None, + help="JSON object merged into request extra_body, e.g. '{\"cfg_text_scale\": 4.0}'.", + ) args = parser.parse_args() @@ -152,6 +175,7 @@ def main(): guidance_scale=args.guidance, seed=args.seed, negative_prompt=args.negative, + extra_body=args.extra_body, ) if image_bytes: diff --git a/examples/online_serving/sensenova_u1/openai_chat_client.py b/examples/online_serving/sensenova_u1/openai_chat_client.py index df9afa83d..50ff202d4 100644 --- a/examples/online_serving/sensenova_u1/openai_chat_client.py +++ b/examples/online_serving/sensenova_u1/openai_chat_client.py @@ -48,7 +48,7 @@ def generate( forwarded as top-level fields in the request payload. The serving layer maps standard keys (height, width, seed, num_inference_steps) to sampling params and forwards model-specific keys (think, cfg_norm, img_cfg_scale, - …) to ``extra_args`` based on the pipeline's ``EXTRA_BODY_PARAMS``. + …) to ``extra_args`` based on the model extra registry. Returns: Image bytes (for image outputs) or text string (for text outputs). @@ -131,7 +131,7 @@ def main(): parser.add_argument("--width", type=int, default=2048, help="Image width") parser.add_argument("--num-steps", type=int, default=50, help="Inference steps") parser.add_argument("--seed", type=int, default=42, help="Random seed") - # Model-specific parameters (forwarded via EXTRA_BODY_PARAMS) + # Model-specific parameters forwarded via the model extra registry. parser.add_argument("--cfg-scale", type=float, default=4.0, help="CFG scale") parser.add_argument("--img-cfg-scale", type=float, help="Image CFG scale (img2img)") parser.add_argument("--cfg-norm", type=str, help="CFG normalization mode") diff --git a/recipes/Bagel/BAGEL-7B-MoT.md b/recipes/Bagel/BAGEL-7B-MoT.md new file mode 100644 index 000000000..a7b0524e4 --- /dev/null +++ b/recipes/Bagel/BAGEL-7B-MoT.md @@ -0,0 +1,227 @@ +# BAGEL-7B-MoT + +> BAGEL image generation through the shared online and offline image examples + +## Summary + +- Vendor: ByteDance Seed +- Model: `ByteDance-Seed/BAGEL-7B-MoT` +- Task: Text-to-image and image-to-image generation +- Mode: Offline inference and OpenAI-compatible online serving +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want to run BAGEL through the shared image examples +instead of model-specific example clients. The generic examples can format +BAGEL text-to-image and image-to-image prompts, select the image output +modality, attach reference images, and forward BAGEL-specific generation +parameters through the pipeline-declared `extra_args` contract. + +## References + +- Upstream model: + [`ByteDance-Seed/BAGEL-7B-MoT`](https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT) +- Related offline example: + [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py) +- Related offline image-to-image example: + [`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_edit.py) +- Related online example: + [`examples/online_serving/text_to_image/openai_chat_client.py`](../../examples/online_serving/text_to_image/openai_chat_client.py) +- Related online image-to-image example: + [`examples/online_serving/image_to_image/openai_chat_client.py`](../../examples/online_serving/image_to_image/openai_chat_client.py) +- Default deploy configs: + [`vllm_omni/deploy/bagel.yaml`](../../vllm_omni/deploy/bagel.yaml), + [`vllm_omni/deploy/bagel_single_stage.yaml`](../../vllm_omni/deploy/bagel_single_stage.yaml) + +## Hardware Support + +This recipe documents the CUDA layouts used by the in-repo BAGEL deploy +configs. The default two-stage config shares one 80 GB GPU; for more headroom, +move the diffusion stage to a second GPU in a custom deploy config. + +## GPU + +### 1x A100 80GB + +#### Environment + +- OS: Linux +- Python: Match the repository requirements for your checkout +- Driver / runtime: NVIDIA CUDA environment with one A100 80 GB GPU +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Offline Commands + +Run text-to-image with the shared offline example from the repository root. +Pick the topology with `--stage-configs-path` and forward BAGEL-specific +generation parameters as a JSON object through `--extra-body`: + +```bash +# Two-stage (Thinker + DiT), shares one GPU by default +python examples/offline_inference/text_to_image/text_to_image.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --stage-configs-path vllm_omni/deploy/bagel.yaml \ + --prompt "A beautiful sunset over mountains" \ + --height 1024 \ + --width 1024 \ + --num-inference-steps 50 \ + --negative-prompt "blurry, low quality" \ + --seed 42 \ + --extra-body '{"timestep_shift": 3.0, "cfg_text_scale": 4.0, "cfg_img_scale": 1.5, "cfg_interval": [0.4, 1.0], "cfg_renorm_type": "global", "cfg_renorm_min": 0.0, "think": false}' \ + --output /tmp/bagel_text2img.png +``` + +```bash +# Single-stage (DiT only, with internal LLM/ViT/VAE) +python examples/offline_inference/text_to_image/text_to_image.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --stage-configs-path vllm_omni/deploy/bagel_single_stage.yaml \ + --prompt "A beautiful sunset over mountains" \ + --height 1024 \ + --width 1024 \ + --num-inference-steps 50 \ + --seed 42 \ + --extra-body '{"timestep_shift": 3.0, "cfg_text_scale": 4.0, "cfg_interval": [0.4, 1.0], "cfg_renorm_type": "global", "cfg_renorm_min": 0.0}' \ + --output /tmp/bagel_text2img_single.png +``` + +The `--extra-body` JSON forwards BAGEL-specific parameters into +`OmniDiffusionSamplingParams.extra_args`. Keys are filtered against the model's +declared `extra_body_params` (see +[`vllm_omni/model_extras/bagel.py`](../../vllm_omni/model_extras/bagel.py)), so +unknown keys for BAGEL are silently dropped. This is the preferred way to pass +model-specific params — `timestep_shift` (use `3.0` for 1024×1024; the +deprecated `--timesteps-shift` flag defaults to `1.0` and is NextStep-oriented), +`cfg_text_scale`, `cfg_img_scale`, `cfg_interval`, `cfg_renorm_type`, +`cfg_renorm_min`, and `think`. `--seed` makes generation reproducible. + +> **Note**: the two-stage `bagel.yaml` runs both stages on GPU 0. If you hit +> CUDA OOM during warmup, lower stage-0 `gpu_memory_utilization` in the YAML or +> split the stages across two GPUs (see the `2x CUDA GPUs` section). + +Run image-to-image with the shared offline image-to-image example: + +```bash +python examples/offline_inference/image_to_image/image_edit.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --prompt "Make the scene look like a watercolor painting" \ + --image /tmp/bagel_text2img.png \ + --height 512 \ + --width 512 \ + --extra-args '{"cfg_text_scale": 4.0, "cfg_img_scale": 1.5}' \ + --output /tmp/bagel_img2img.png +``` + +The `--extra-args` JSON forwards BAGEL-specific parameters (e.g. `cfg_text_scale`, +`cfg_img_scale`, `cfg_interval`, `cfg_renorm_type`) into +`OmniDiffusionSamplingParams.extra_args` via the model-extras registry. + +#### Online Commands + +Start the OpenAI-compatible server: + +```bash +vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 +``` + +To use the single-stage topology online: + +```bash +vllm serve ByteDance-Seed/BAGEL-7B-MoT \ + --omni \ + --port 8091 \ + --deploy-config vllm_omni/deploy/bagel_single_stage.yaml +``` + +Send a text-to-image request with BAGEL-specific generation parameters: + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "<|im_start|>A beautiful sunset over mountains<|im_end|>"} + ] + } + ], + "modalities": ["image"], + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 50, + "cfg_text_scale": 4.0, + "cfg_img_scale": 1.5, + "negative_prompt": "blurry, low quality", + "seed": 42 + } + }' +``` + +The important part is that BAGEL-specific keys such as `cfg_text_scale`, +`cfg_img_scale`, `cfg_interval`, `cfg_renorm_type`, and `cfg_renorm_min` belong +in `extra_body`. The serving layer routes the declared keys into +`OmniDiffusionSamplingParams.extra_args` for the BAGEL pipeline. + +#### Verification + +Decode the returned data URL into an image: + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "<|im_start|>A ceramic teapot on a wooden table<|im_end|>"} + ] + } + ], + "modalities": ["image"], + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 25, + "cfg_text_scale": 4.0, + "seed": 42 + } + }' | jq -r '.choices[]?.message.content[]? | select(.image_url.url) | .image_url.url' \ + | head -n 1 \ + | cut -d',' -f2- \ + | base64 -d > /tmp/bagel_online.png + +ls -lh /tmp/bagel_online.png +``` + +### 2x CUDA GPUs + +Create a custom deploy config from `vllm_omni/deploy/bagel.yaml` and move the +diffusion stage to GPU 1: + +```yaml +stages: + - stage_id: 0 + devices: "0" + # keep the remaining stage-0 settings from bagel.yaml + - stage_id: 1 + devices: "1" + # keep the remaining stage-1 settings from bagel.yaml +``` + +Then start serving with that config: + +```bash +vllm serve ByteDance-Seed/BAGEL-7B-MoT \ + --omni \ + --port 8091 \ + --deploy-config /path/to/custom_bagel_2gpu.yaml +``` + +Use the online curl request from the `1x A100 80GB` section to verify that the +server returns an image. diff --git a/recipes/README.md b/recipes/README.md index a1c6973d5..1f5b5be57 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -29,6 +29,7 @@ recipes/ |--------|------|----------| | [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB | | [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB | +| [`Bagel/BAGEL-7B-MoT.md`](./Bagel/BAGEL-7B-MoT.md) | Text-to-image with shared online/offline examples | 1x A100 80GB / 2x CUDA GPUs | | [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB | | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB | | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 | diff --git a/recipes/SenseNova/SenseNova-U1.md b/recipes/SenseNova/SenseNova-U1.md index b1a14eaa2..5cb30a21d 100644 --- a/recipes/SenseNova/SenseNova-U1.md +++ b/recipes/SenseNova/SenseNova-U1.md @@ -72,6 +72,31 @@ pytest -s -v tests/e2e/offline_inference/test_sensenova_u1_text2img.py \ - No deploy YAML needed — the engine auto-generates a single-stage diffusion config. - Think mode (`--think`) is recommended for higher image quality. +#### Text-to-Image via the shared example + +SenseNova-U1 is also registered with the shared text-to-image example, so you +can run it through `text_to_image.py` instead of `end2end.py`. Forward the +SenseNova-specific generation parameters as a JSON object through `--extra-body`: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model SenseNova/SenseNova-U1-8B-MoT \ + --prompt "A beautiful sunset over mountains" \ + --width 2048 --height 2048 \ + --num-inference-steps 50 \ + --seed 42 \ + --extra-body '{"think": true, "cfg_scale": 4.0, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \ + --output sensenova_text2img.png +``` + +The `--extra-body` keys are filtered against the model's declared +`extra_body_params` (see +[`vllm_omni/model_extras/sensenova_u1.py`](../../vllm_omni/model_extras/sensenova_u1.py)): +`think`, `cfg_scale`, `cfg_norm`, `timestep_shift`, `t_eps`, `img_cfg_scale`, +and `max_tokens`. No deploy YAML is needed — the engine auto-generates a +single-stage diffusion config. For richer task modalities (img2text, text2text) +and CPU offload, use the dedicated `end2end.py` script described above. + #### Image-to-Image Editing (img2img) ```bash diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py new file mode 100644 index 000000000..09c732d07 --- /dev/null +++ b/tests/model_extras/test_model_extras.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import pytest +from PIL import Image + +from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_extras import ( + build_image_to_image_prompt, + build_text_to_image_prompt, + get_extra_body_params, + get_extra_output_params, + should_init_extra_args_for_non_diffusion_stages, +) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_bagel_extra_registry_declares_request_and_response_params() -> None: + assert get_extra_body_params("BagelPipeline") == frozenset( + { + "cfg_text_scale", + "cfg_img_scale", + "cfg_interval", + "cfg_renorm_type", + "cfg_renorm_min", + "negative_prompt", + "think", + "max_think_tokens", + "do_sample", + "text_temperature", + "timestep_shift", + } + ) + assert get_extra_output_params("BagelPipeline") == frozenset({"text_output", "think_text"}) + assert should_init_extra_args_for_non_diffusion_stages("BagelPipeline") is True + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_sensenova_extra_registry_declares_request_and_response_params() -> None: + assert get_extra_body_params("SenseNovaU1Pipeline") == frozenset( + { + "think", + "cfg_scale", + "cfg_norm", + "timestep_shift", + "t_eps", + "img_cfg_scale", + "max_tokens", + } + ) + assert get_extra_output_params("SenseNovaU1Pipeline") == frozenset({"think_text"}) + assert should_init_extra_args_for_non_diffusion_stages("SenseNovaU1Pipeline") is False + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_unknown_pipeline_has_empty_extra_registry() -> None: + assert get_extra_body_params("UnknownPipeline") == frozenset() + assert get_extra_output_params("UnknownPipeline") == frozenset() + assert should_init_extra_args_for_non_diffusion_stages("UnknownPipeline") is False + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_bagel_text_to_image_prompt_builder() -> None: + assert build_text_to_image_prompt( + "BagelPipeline", + prompt="a cat", + negative_prompt="blurry", + height=512, + width=768, + ) == { + "prompt": "<|im_start|>a cat<|im_end|>", + "modalities": ["image"], + "mm_processor_kwargs": { + "target_h": 512, + "target_w": 768, + "modalities": ["image"], + }, + "negative_prompt": "blurry", + } + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_bagel_image_to_image_prompt_builder() -> None: + dummy_image = Image.new("RGB", (64, 64)) + result = build_image_to_image_prompt( + "BagelPipeline", + prompt="paint it", + negative_prompt="ugly", + input_image=dummy_image, + height=256, + width=256, + ) + assert result["prompt"] == "<|fim_middle|><|im_start|>paint it<|im_end|>" + assert result["modalities"] == ["img2img"] + assert result["multi_modal_data"]["img2img"] is dummy_image + assert result["mm_processor_kwargs"]["target_h"] == 256 + assert result["mm_processor_kwargs"]["target_w"] == 256 + assert result["negative_prompt"] == "ugly" + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_unknown_pipeline_uses_default_text_to_image_prompt() -> None: + assert build_text_to_image_prompt( + "UnknownPipeline", + prompt="a cat", + negative_prompt=None, + height=512, + width=512, + ) == {"prompt": "a cat"} + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_unknown_pipeline_uses_default_image_to_image_prompt() -> None: + dummy_image = Image.new("RGB", (64, 64)) + result = build_image_to_image_prompt( + "UnknownPipeline", + prompt="edit", + negative_prompt=None, + input_image=dummy_image, + ) + assert result == { + "prompt": "edit", + "multi_modal_data": {"image": dummy_image}, + } + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_declared_extra_args_apply_to_existing_sampling_params() -> None: + params = OmniDiffusionSamplingParams(extra_args={"existing": 1}) + + declared_extra_params: frozenset[str] = frozenset({"cfg_text_scale", "think"}) + apply_declared_extra_args( + params, + declared_extra_params, + { + "cfg_text_scale": 4.0, + "think": False, + "unknown": "ignored", + }, + ) + + assert params.extra_args == { + "existing": 1, + "cfg_text_scale": 4.0, + "think": False, + } diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index a2dfb5a15..c4cca97bb 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -423,6 +423,67 @@ class DiffusionCacheConfig: raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{item}'") +def resolve_model_class_name(model: str | None, diffusion_load_format: str = "default") -> str | None: + """Resolve the diffusion pipeline class name from the model config. + + Read-only counterpart of ``OmniDiffusionConfig.enrich_config``, safe to call + client-side. Returns ``None`` if the pipeline can't be determined. + """ + from vllm.transformers_utils.config import get_hf_file_to_dict + + if not model: + return None + + is_lance_subfolder = os.path.basename(str(model).rstrip("/")) in {"Lance_3B", "Lance_3B_Video"} + + # Diffusers models: read _class_name from model_index.json. + try: + model_index = get_hf_file_to_dict("model_index.json", model) + except Exception: + model_index = None + if model_index is not None: + return model_index.get("_class_name") + if diffusion_load_format == "diffusers": + return "DiffusersAdapterPipeline" + + # Other models: map model_type / architecture from config.json. + try: + cfg = get_hf_file_to_dict("config.json", model) or {} + except Exception: + cfg = {} + model_type = cfg.get("model_type") + architectures = cfg.get("architectures") or [] + + if model_type == "bagel" or "BagelForConditionalGeneration" in architectures: + return "BagelPipeline" + if ( + model_type == "lance" + or "LancePipeline" in architectures + or cfg.get("model_name") == "Lance" + or is_lance_subfolder + ): + return "LancePipeline" + if model_type == "neo_chat": + return "SenseNovaU1Pipeline" + if "BailingMM2NativeForConditionalGeneration" in architectures or model_type in ( + "bailingmm_moe_v2_lite", + "ming_flash_omni", + "ming_flash_omni_thinker", + ): + return "MingImagePipeline" + if model_type == "nextstep": + return "NextStep11Pipeline" + if model_type == "s2v": + return "WanS2VPipeline" + if model_type == "vla": + from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero + + return "DreamZeroPipeline" if _looks_like_dreamzero(model) else None + if len(architectures) == 1: + return architectures[0] + return None + + @dataclass class OmniDiffusionConfig: # Model and path configuration (for convenience) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 3b4843af0..a7869b140 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -27,8 +27,6 @@ from vllm_omni.diffusion.data import ( from vllm_omni.diffusion.executor.abstract import DiffusionExecutor from vllm_omni.diffusion.io_support import ( get_dummy_run_num_frames, - get_extra_body_params, - get_extra_output_params, image_color_format, supports_audio_output, supports_multimodal_input, @@ -61,8 +59,6 @@ __all__ = [ "_RpcTask", "_move_tensor_tree_to_cpu", "get_dummy_run_num_frames", - "get_extra_body_params", - "get_extra_output_params", "image_color_format", "supports_audio_output", "supports_multimodal_input", diff --git a/vllm_omni/diffusion/io_support.py b/vllm_omni/diffusion/io_support.py index 9c7b2b027..8d1405a6c 100644 --- a/vllm_omni/diffusion/io_support.py +++ b/vllm_omni/diffusion/io_support.py @@ -47,31 +47,3 @@ def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool) if model_cls is not None and hasattr(model_cls, "dummy_run_num_frames"): return int(getattr(model_cls, "dummy_run_num_frames")) return 2 if supports_audio_input or supports_audio_output(model_class_name) else 1 - - -def get_extra_body_params(model_class_name: str) -> frozenset[str]: - """Return the set of extra_body keys accepted by a pipeline. - - Each pipeline can declare ``EXTRA_BODY_PARAMS: ClassVar[frozenset[str]]`` - to advertise which request-level parameters should be forwarded from - ``extra_body`` to ``OmniDiffusionSamplingParams.extra_args``. - Returns an empty frozenset when the pipeline does not declare any. - """ - model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) - if model_cls is None: - return frozenset() - return frozenset(getattr(model_cls, "EXTRA_BODY_PARAMS", frozenset())) - - -def get_extra_output_params(model_class_name: str) -> frozenset[str]: - """Return the set of custom_output keys to expose in API response metrics. - - Each pipeline can declare ``EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]]`` - to advertise which ``DiffusionOutput.custom_output`` keys should be - copied into the response ``metrics`` dict. - Returns an empty frozenset when the pipeline does not declare any. - """ - model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) - if model_cls is None: - return frozenset() - return frozenset(getattr(model_cls, "EXTRA_OUTPUT_PARAMS", frozenset())) diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py index fa7f95c66..44446ae83 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py @@ -504,30 +504,6 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli support_image_input = True - # Model-specific parameters accepted via ``extra_body`` in the online - # serving API. The serving layer reads this set and forwards matching - # keys from the request payload to ``OmniDiffusionSamplingParams.extra_args``. - EXTRA_BODY_PARAMS: ClassVar[frozenset[str]] = frozenset( - { - "think", - "cfg_scale", - "cfg_norm", - "timestep_shift", - "t_eps", - "img_cfg_scale", - "max_tokens", - } - ) - - # Keys from ``DiffusionOutput.custom_output`` that should be surfaced in - # the API response ``metrics`` dict. The serving layer reads this set and - # copies matching entries from ``custom_output`` into the response. - EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]] = frozenset( - { - "think_text", - } - ) - # CPU-offload protocol: language_model carries the denoising blocks; the # vision and FM modules are lightweight encoders pinned on GPU during the # diffusion loop. There is no separate VAE. diff --git a/vllm_omni/diffusion/utils/param_utils.py b/vllm_omni/diffusion/utils/param_utils.py new file mode 100644 index 000000000..e56fcd072 --- /dev/null +++ b/vllm_omni/diffusion/utils/param_utils.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + +def apply_declared_extra_args( + sampling_params: OmniDiffusionSamplingParams, + declared_params: frozenset[str], + user_kwargs: dict[str, object], +) -> None: + """Route pipeline-declared request params into ``sampling_params.extra_args``. + + Both online serving and offline examples call this so that model-specific + keys (e.g. ``cfg_text_scale`` for BAGEL) end up in ``extra_args`` instead + of being silently dropped. + + This is a no-op when no declared params are present in ``user_kwargs``, so + it is safe to call on non-diffusion (e.g. AR) sampling params whose + ``extra_args`` defaults to ``None``. + """ + declared = {key: user_kwargs[key] for key in declared_params if user_kwargs.get(key) is not None} + if not declared: + return + sampling_params.extra_args = {**(sampling_params.extra_args or {}), **declared} diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index dc1385e38..721e58c7b 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -220,6 +220,8 @@ class AsyncOmniEngine: self.model = model self.tokenizer = tokenizer self.diffusion_batch_size = diffusion_batch_size + # Cached by get_diffusion_od_config(). + self._diffusion_od_config_view: Any = None startup_timeout = int(init_timeout) # Forwarded into Orchestrator so its _forward_to_next_stage path can # emit per-edge transfer_tx_s / transfer_size_bytes histograms. @@ -323,6 +325,20 @@ class AsyncOmniEngine: logger.info(f"[AsyncOmniEngine] Orchestrator ready with {self.num_stages} stages") + def get_diffusion_od_config(self) -> Any: + """Expose the diffusion ``model_class_name`` to client-side model-extras. + + The worker holds the full config; here we just resolve the pipeline class + name from the model config (cached). ``model_class_name`` may be ``None``. + """ + if self._diffusion_od_config_view is None: + from types import SimpleNamespace + + from vllm_omni.diffusion.data import resolve_model_class_name + + self._diffusion_od_config_view = SimpleNamespace(model_class_name=resolve_model_class_name(self.model)) + return self._diffusion_od_config_view + def _initialize_stages(self, stage_init_timeout: int) -> None: """Initialize stage clients/processors via StageRuntime and assign to self.""" self._runtime = create_stage_runtime( diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 24998436e..9a0a9ed1c 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -23,7 +23,7 @@ from vllm.entrypoints.chat_utils import ( make_tool_call_id, ) -from vllm_omni.diffusion.io_support import get_extra_body_params, get_extra_output_params +from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse from vllm_omni.entrypoints.utils import coerce_param_message_types @@ -33,6 +33,7 @@ from vllm_omni.metrics.modality import ( observe_audio_first_packet, observe_audio_streaming_finalize, ) +from vllm_omni.model_extras import get_extra_body_params, get_extra_output_params try: import soundfile @@ -162,12 +163,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): instance._diffusion_mode = True instance._diffusion_engine = diffusion_engine instance._diffusion_model_name = model_name + instance._diffusion_extra_body_params = None + instance._diffusion_extra_output_params = None instance.engine_client = None instance.has_kv_connector = False + # Extra body/output params are resolved lazily on first use; see + # _get_diffusion_extra_body_params / _get_diffusion_extra_output_params. return instance def _get_diffusion_extra_body_params(self) -> frozenset[str]: - """Return the pipeline-declared ``EXTRA_BODY_PARAMS`` (cached).""" + """Return model-specific extra_body params from the extra registry.""" if self._diffusion_extra_body_params is not None: return self._diffusion_extra_body_params @@ -177,7 +182,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if od_config is not None and getattr(od_config, "model_class_name", None): params = get_extra_body_params(od_config.model_class_name) except Exception as e: - logger.warning("Failed to read EXTRA_BODY_PARAMS from pipeline: %s", e) + logger.warning("Failed to read model extra_body params: %s", e) self._diffusion_extra_body_params = params return params @@ -186,7 +191,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): self, custom_output: dict[str, Any] | None, ) -> dict[str, Any] | None: - """Pick pipeline-declared ``EXTRA_OUTPUT_PARAMS`` keys from *custom_output*.""" + """Pick model-specific extra output keys from *custom_output*.""" if not custom_output: return None @@ -197,7 +202,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if od_config is not None and getattr(od_config, "model_class_name", None): params = get_extra_output_params(od_config.model_class_name) except Exception as e: - logger.warning("Failed to read EXTRA_OUTPUT_PARAMS from pipeline: %s", e) + logger.warning("Failed to read model extra output params: %s", e) self._diffusion_extra_output_params = params if not self._diffusion_extra_output_params: @@ -397,8 +402,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): request.modalities = output_modalities if output_modalities is not None else engine_output_modalities num_inference_steps = None - cfg_text_scale = None - cfg_img_scale = None + extra_body: dict[str, Any] = {} # Omni multistage image generation: Stage-0 (AR) should receive a clean # text prompt (and optional conditioning image/size) so the model's own # processor can construct the correct inputs. @@ -416,9 +420,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # `extra_body` is flattented and merged into the payload's root. # These extra fields are accessible via `model_extra` property (from Pydantic base class). # When sending raw request with curl, no flattening happens. Directly read the `extra_body` dict. - extra_body = getattr(request, "extra_body", None) - if not extra_body: - extra_body = request.model_extra or {} + extra_body = getattr(request, "extra_body", None) or request.model_extra or {} height, width = self._resolve_height_width_from_extra_body(extra_body) @@ -430,9 +432,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): num_inference_steps = None negative_prompt = extra_body.get("negative_prompt") - cfg_text_scale = extra_body.get("cfg_text_scale") - cfg_img_scale = extra_body.get("cfg_img_scale") - engine_prompt_image: dict[str, Any] | None = None if reference_images: # Best-effort decode first reference image for i2i. @@ -559,15 +558,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): sp.width = _image_gen_width if hasattr(sp, "num_inference_steps") and num_inference_steps is not None: sp.num_inference_steps = num_inference_steps - if hasattr(sp, "extra_args") and sp.extra_args is not None: - if cfg_text_scale is not None: - sp.extra_args["cfg_text_scale"] = cfg_text_scale - if cfg_img_scale is not None: - sp.extra_args["cfg_img_scale"] = cfg_img_scale - for _key in self._get_diffusion_extra_body_params(): - _val = extra_body.get(_key) - if _val is not None: - sp.extra_args[_key] = _val + apply_declared_extra_args(sp, self._get_diffusion_extra_body_params(), extra_body) self._log_inputs( request_id, @@ -2596,6 +2587,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): layers=layers, resolution=resolution, ) + apply_declared_extra_args( + default_stage_params, + self._get_diffusion_extra_body_params(), + extra_body, + ) if lora_body and isinstance(lora_body, dict): try: lora_req, lora_scale = parse_lora_request(lora_body) @@ -2990,8 +2986,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): num_inference_steps = extra_body.get("num_inference_steps") guidance_scale = extra_body.get("guidance_scale") true_cfg_scale = extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale") - cfg_text_scale = extra_body.get("cfg_text_scale") - cfg_img_scale = extra_body.get("cfg_img_scale") seed = extra_body.get("seed") if seed is None: seed = getattr(request, "seed", None) @@ -3049,14 +3043,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): gen_params.guidance_scale = guidance_scale if true_cfg_scale is not None: gen_params.true_cfg_scale = true_cfg_scale - if cfg_text_scale is not None: - gen_params.extra_args["cfg_text_scale"] = cfg_text_scale - if cfg_img_scale is not None: - gen_params.extra_args["cfg_img_scale"] = cfg_img_scale - for _key in self._get_diffusion_extra_body_params(): - _val = extra_body.get(_key) - if _val is not None: - gen_params.extra_args[_key] = _val + apply_declared_extra_args(gen_params, self._get_diffusion_extra_body_params(), extra_body) if num_frames is not None: gen_params.num_frames = num_frames if guidance_scale_2 is not None: diff --git a/vllm_omni/model_extras/__init__.py b/vllm_omni/model_extras/__init__.py new file mode 100644 index 000000000..e2f98cd41 --- /dev/null +++ b/vllm_omni/model_extras/__init__.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_omni.model_extras.registry import ( + build_image_to_image_prompt, + build_text_to_image_prompt, + get_extra_body_params, + get_extra_output_params, + get_model_class_name, + should_init_extra_args_for_non_diffusion_stages, +) + +__all__ = [ + "build_image_to_image_prompt", + "build_text_to_image_prompt", + "get_extra_body_params", + "get_extra_output_params", + "get_model_class_name", + "should_init_extra_args_for_non_diffusion_stages", +] diff --git a/vllm_omni/model_extras/bagel.py b/vllm_omni/model_extras/bagel.py new file mode 100644 index 000000000..e8cbcbf11 --- /dev/null +++ b/vllm_omni/model_extras/bagel.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from typing import Any + +from PIL import Image + + +def build_text_to_image_prompt( + prompt: str, + negative_prompt: str | None, + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + text_prompt: dict[str, Any] = { + "prompt": f"<|im_start|>{prompt}<|im_end|>", + "modalities": ["image"], + "mm_processor_kwargs": { + "target_h": height, + "target_w": width, + "modalities": ["image"], + }, + } + if negative_prompt is not None: + text_prompt["negative_prompt"] = negative_prompt + return text_prompt + + +def build_image_to_image_prompt( + prompt: str, + negative_prompt: str | None, + input_image: Image.Image | list[Image.Image], + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + img_prompt: dict[str, Any] = { + "prompt": f"<|fim_middle|><|im_start|>{prompt}<|im_end|>", + "modalities": ["img2img"], + "multi_modal_data": {"img2img": input_image}, + "mm_processor_kwargs": {"modalities": ["img2img"]}, + } + if height is not None: + img_prompt["mm_processor_kwargs"]["target_h"] = height + if width is not None: + img_prompt["mm_processor_kwargs"]["target_w"] = width + if negative_prompt is not None: + img_prompt["negative_prompt"] = negative_prompt + return img_prompt + + +BAGEL_EXTRA_BODY_PARAMS = frozenset( + { + "cfg_text_scale", + "cfg_img_scale", + "cfg_interval", + "cfg_renorm_type", + "cfg_renorm_min", + "negative_prompt", + "think", + "max_think_tokens", + "do_sample", + "text_temperature", + "timestep_shift", + } +) +BAGEL_EXTRA_OUTPUT_PARAMS = frozenset( + { + "text_output", + "think_text", + } +) +BAGEL_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py new file mode 100644 index 000000000..639a4b67d --- /dev/null +++ b/vllm_omni/model_extras/registry.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +from PIL import Image + +from vllm_omni.model_extras.bagel import ( + BAGEL_EXTRA_BODY_PARAMS, + BAGEL_EXTRA_OUTPUT_PARAMS, + BAGEL_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, +) +from vllm_omni.model_extras.bagel import ( + build_image_to_image_prompt as build_bagel_image_to_image_prompt, +) +from vllm_omni.model_extras.bagel import ( + build_text_to_image_prompt as build_bagel_text_to_image_prompt, +) +from vllm_omni.model_extras.sensenova_u1 import ( + SENSENOVA_U1_EXTRA_BODY_PARAMS, + SENSENOVA_U1_EXTRA_OUTPUT_PARAMS, +) + +# TODO: extend with TextToVideoPromptBuilder, ImageToVideoPromptBuilder +# once a video model needs model-specific prompt construction (follow-up PR). +TextToImagePromptBuilder = Callable[ + [str, str | None, int | None, int | None], + dict[str, Any], +] +ImageToImagePromptBuilder = Callable[ + [str, str | None, "Image.Image | list[Image.Image]", int | None, int | None], + dict[str, Any], +] + + +def default_text_to_image_prompt( + prompt: str, + negative_prompt: str | None, + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + result: dict[str, Any] = {"prompt": prompt} + if negative_prompt is not None: + result["negative_prompt"] = negative_prompt + return result + + +def default_image_to_image_prompt( + prompt: str, + negative_prompt: str | None, + input_image: Image.Image | list[Image.Image], + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + result: dict[str, Any] = { + "prompt": prompt, + "multi_modal_data": {"image": input_image}, + } + if negative_prompt is not None: + result["negative_prompt"] = negative_prompt + return result + + +_EXTRA_SPECS: dict[str, dict[str, Any]] = { + "BagelPipeline": { + "extra_body_params": BAGEL_EXTRA_BODY_PARAMS, + "extra_output_params": BAGEL_EXTRA_OUTPUT_PARAMS, + "init_extra_args_for_non_diffusion_stages": BAGEL_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, + "text_to_image_prompt_builder": build_bagel_text_to_image_prompt, + "image_to_image_prompt_builder": build_bagel_image_to_image_prompt, + }, + "SenseNovaU1Pipeline": { + "extra_body_params": SENSENOVA_U1_EXTRA_BODY_PARAMS, + "extra_output_params": SENSENOVA_U1_EXTRA_OUTPUT_PARAMS, + }, +} + + +def _get_spec(model_class_name: str | None) -> dict[str, Any] | None: + if not model_class_name: + return None + return _EXTRA_SPECS.get(model_class_name) + + +def get_model_class_name(omni: Any) -> str | None: + """Extract model_class_name from an Omni/AsyncOmni instance. + + This hides the internal ODConfig plumbing from example scripts. + """ + engine = getattr(omni, "engine", None) + if engine is None: + return None + od_config = getattr(engine, "od_config", None) + if od_config is None and hasattr(engine, "get_diffusion_od_config"): + od_config = engine.get_diffusion_od_config() + return getattr(od_config, "model_class_name", None) if od_config else None + + +def get_extra_body_params(model_class_name: str | None) -> frozenset[str]: + spec = _get_spec(model_class_name) + return spec.get("extra_body_params", frozenset()) if spec is not None else frozenset() + + +def get_extra_output_params(model_class_name: str | None) -> frozenset[str]: + spec = _get_spec(model_class_name) + return spec.get("extra_output_params", frozenset()) if spec is not None else frozenset() + + +def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None) -> bool: + spec = _get_spec(model_class_name) + return bool(spec and spec.get("init_extra_args_for_non_diffusion_stages", False)) + + +def build_text_to_image_prompt( + model_class_name: str | None, + prompt: str, + negative_prompt: str | None, + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + spec = _get_spec(model_class_name) + builder: TextToImagePromptBuilder = ( + spec.get("text_to_image_prompt_builder", default_text_to_image_prompt) + if spec is not None + else default_text_to_image_prompt + ) + return builder(prompt, negative_prompt, height, width) + + +def build_image_to_image_prompt( + model_class_name: str | None, + prompt: str, + negative_prompt: str | None, + input_image: Image.Image | list[Image.Image], + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + spec = _get_spec(model_class_name) + builder: ImageToImagePromptBuilder = ( + spec.get("image_to_image_prompt_builder", default_image_to_image_prompt) + if spec is not None + else default_image_to_image_prompt + ) + return builder(prompt, negative_prompt, input_image, height, width) diff --git a/vllm_omni/model_extras/sensenova_u1.py b/vllm_omni/model_extras/sensenova_u1.py new file mode 100644 index 000000000..29fcc52c4 --- /dev/null +++ b/vllm_omni/model_extras/sensenova_u1.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +SENSENOVA_U1_EXTRA_BODY_PARAMS = frozenset( + { + "think", + "cfg_scale", + "cfg_norm", + "timestep_shift", + "t_eps", + "img_cfg_scale", + "max_tokens", + } +) +SENSENOVA_U1_EXTRA_OUTPUT_PARAMS = frozenset( + { + "think_text", + } +) From 0383fde38a9b2940ade4b479666f84040fcdfe03 Mon Sep 17 00:00:00 2001 From: Zening Chen <59945561+zeningc@users.noreply.github.com> Date: Thu, 18 Jun 2026 02:31:36 -0700 Subject: [PATCH 06/40] [Doc] Qwen3-TTS: add 0.6B on 1x RTX 4090 and fix broken links (#4026) Signed-off-by: Zening Chen <59945561+zeningc@users.noreply.github.com> Signed-off-by: zeningc --- recipes/Qwen/Qwen3-TTS.md | 103 ++++++++++++++++++++++++++++++++++---- 1 file changed, 93 insertions(+), 10 deletions(-) diff --git a/recipes/Qwen/Qwen3-TTS.md b/recipes/Qwen/Qwen3-TTS.md index 73a14a630..164fba2c9 100644 --- a/recipes/Qwen/Qwen3-TTS.md +++ b/recipes/Qwen/Qwen3-TTS.md @@ -28,11 +28,9 @@ Smaller 0.6B variants are also available for `CustomVoice` and `Base`. ## References -- Upstream or canonical docs: - [`docs/user_guide/examples/online_serving/qwen3_tts.md`](../../docs/user_guide/examples/online_serving/qwen3_tts.md) - Related examples under `examples/`: - [`examples/online_serving/qwen3_tts/`](../../examples/online_serving/qwen3_tts/), - [`examples/offline_inference/qwen3_tts/`](../../examples/offline_inference/qwen3_tts/) + [`examples/online_serving/text_to_speech/qwen3_tts/`](../../examples/online_serving/text_to_speech/qwen3_tts/), + [`examples/offline_inference/text_to_speech/qwen3_tts/`](../../examples/offline_inference/text_to_speech/qwen3_tts/) - Related issue or discussion: [RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645) @@ -67,9 +65,9 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \ Alternatively, use the convenience script: ```bash -./examples/online_serving/qwen3_tts/run_server.sh # Default: CustomVoice -./examples/online_serving/qwen3_tts/run_server.sh VoiceDesign # VoiceDesign -./examples/online_serving/qwen3_tts/run_server.sh Base # Base (voice clone) +./examples/online_serving/text_to_speech/qwen3_tts/run_server.sh # Default: CustomVoice +./examples/online_serving/text_to_speech/qwen3_tts/run_server.sh VoiceDesign # VoiceDesign +./examples/online_serving/text_to_speech/qwen3_tts/run_server.sh Base # Base (voice clone) ``` The bundled deploy config (`vllm_omni/deploy/qwen3_tts.yaml`) enables async @@ -117,7 +115,7 @@ curl http://localhost:8091/v1/audio/voices **Using the Python client:** ```bash -cd examples/online_serving/qwen3_tts +cd examples/online_serving/text_to_speech/qwen3_tts # CustomVoice python openai_speech_client.py \ @@ -157,8 +155,8 @@ curl -X POST http://localhost:8091/v1/audio/speech \ **Offline inference (no server needed):** ```bash -python examples/offline_inference/qwen3_tts/end2end.py --query-type CustomVoice -python examples/offline_inference/qwen3_tts/end2end.py --query-type CustomVoice --streaming +python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-type CustomVoice +python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-type CustomVoice --streaming ``` ## Notes @@ -168,3 +166,88 @@ python examples/offline_inference/qwen3_tts/end2end.py --query-type CustomVoice - Async chunking: Enabled by default in `qwen3_tts.yaml` for streaming-friendly first-audio latency. Streaming requires `stream=true` with `response_format="pcm"`. - Task/model matching: Each task type requires its matching model checkpoint. Using a CustomVoice model for a Base (voice clone) request will fail. - Known limitations: The server serves one model variant at a time. To switch task types (e.g., CustomVoice to Base), restart the server with the corresponding model. + +## Hardware Support + +## GPU + +### 1x RTX 4090 24GB (0.6B CustomVoice) + +#### Environment + +- OS: Ubuntu 24.04 LTS +- Python: 3.12.3 +- PyTorch: 2.11.0+cu130 +- Driver / runtime: NVIDIA 580.126.09 / CUDA 13.0 +- GPU: NVIDIA GeForce RTX 4090, 24 GB +- vLLM version: 0.21.0 +- vLLM-Omni version or commit: 0.21.0rc3.dev91+gd4c13950 + +#### Command + +```bash +vllm serve Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \ + --deploy-config vllm_omni/deploy/qwen3_tts.yaml \ + --omni --port 8091 +``` + +The default deploy config (`qwen3_tts.yaml`) works without modification on the +RTX 4090. Both stages (talker + code2wav) share GPU 0 with +`gpu_memory_utilization: 0.3` each. + +#### Verification + +**English synthesis:** + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, this is Qwen3-TTS running on RTX 4090.", + "voice": "vivian", + "language": "English" + }' --output test_english.wav +``` + +**Chinese synthesis:** + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "你好，这是在RTX 4090上运行的语音合成测试。", + "voice": "vivian", + "language": "Chinese" + }' --output test_chinese.wav +``` + +**With emotion instruction:** + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "I am so excited about this!", + "voice": "vivian", + "language": "English", + "instructions": "Speak with great enthusiasm" + }' --output test_emotion.wav +``` + +#### Notes + +- Memory usage: **~13.5 GiB / 24 GiB** at idle with the default deploy config + (`gpu_memory_utilization: 0.3` per stage). The 0.6B model weights occupy only + ~2.4 GiB (Stage 0: 1.91 GiB, Stage 1: 0.45 GiB); the remainder is KV cache + pre-allocated at startup. To reduce idle footprint to ~5 GiB, use a custom + deploy config with lower utilization (inference peak ~10 GiB): + + ```yaml + # Custom deploy config for Qwen3-TTS-12Hz-0.6B on RTX 4090 + # Copy vllm_omni/deploy/qwen3_tts.yaml and override gpu_memory_utilization: + stages: + - stage_id: 0 + gpu_memory_utilization: 0.15 + - stage_id: 1 + gpu_memory_utilization: 0.15 + ``` From 2d2258890b1a3459e1a6b14a0d2255398dac1ae8 Mon Sep 17 00:00:00 2001 From: zhumingjue138 Date: Thu, 18 Jun 2026 17:53:34 +0800 Subject: [PATCH 07/40] [CI] skip ci for issue 4537 (#4538) Signed-off-by: zhumingjue --- tests/diffusion/offloader/test_diffusion_cpu_offload.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/diffusion/offloader/test_diffusion_cpu_offload.py b/tests/diffusion/offloader/test_diffusion_cpu_offload.py index 7d155fef9..22f6839fd 100644 --- a/tests/diffusion/offloader/test_diffusion_cpu_offload.py +++ b/tests/diffusion/offloader/test_diffusion_cpu_offload.py @@ -102,6 +102,8 @@ def check_audio_determinism(audio1, audio2, atol=1e-2): @hardware_test(res={"cuda": "L4", "rocm": "MI325"}) @pytest.mark.parametrize("model_name", list(MODELS.keys())) def test_cpu_offload_diffusion_model(model_name: str): + if model_name == "OmniGen2/OmniGen2": + pytest.skip("issue #4537") if model_name in _GATED_MODELS: _skip_if_gated_repo_inaccessible(model_name) try: From 8f28b969e91e00db23ad24e687604f2695a67255 Mon Sep 17 00:00:00 2001 From: Honghan Zhu Date: Thu, 18 Jun 2026 18:57:04 +0800 Subject: [PATCH 08/40] [Feat] DreamZero fix tp & cross attn cache (#4154) --- tests/dreamzero/test_crossattn_cache.py | 114 ++++++++++++++++ tests/dreamzero/test_fused_qk_rms_norm.py | 112 +++++++++++++++ tests/dreamzero/test_qkv_fusion.py | 117 ++++++++++++++++ .../models/dreamzero/causal_wan_model.py | 127 +++++++++++++++--- .../models/dreamzero/pipeline_dreamzero.py | 21 ++- .../models/dreamzero/state_dreamzero.py | 8 +- 6 files changed, 479 insertions(+), 20 deletions(-) create mode 100644 tests/dreamzero/test_crossattn_cache.py create mode 100644 tests/dreamzero/test_fused_qk_rms_norm.py create mode 100644 tests/dreamzero/test_qkv_fusion.py diff --git a/tests/dreamzero/test_crossattn_cache.py b/tests/dreamzero/test_crossattn_cache.py new file mode 100644 index 000000000..94361c090 --- /dev/null +++ b/tests/dreamzero/test_crossattn_cache.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Tests that DreamZero's cross-attention KV cache is actually init'd and reused. + +DreamZero's I2V cross-attention caches the text k/v and image k_img/v_img in a +per-layer ``crossattn_cache`` dict: they are computed on the first forward of a +session (``is_init`` False -> True) and reused on every later step, because the +text/image context is session-invariant. Only the query (which depends on the +per-step hidden state) is recomputed. + +These tests lock in that ``WanI2VCrossAttention.forward`` initialises the cache +on the first call and reuses it (not recompute) on later calls, and recomputes +every call when no cache is provided. + +CPU-only: ``WanI2VCrossAttention.__init__`` needs a tensor-parallel group (it +builds ColumnParallel/RowParallel linears), so we bypass ``__init__`` and stub +the projections/attention to count calls and exercise only the caching logic. +""" + +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm_omni.diffusion.models.dreamzero.causal_wan_model import WanI2VCrossAttention + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_CLIP_LEN = 257 # image tokens the I2V cross-attn splits off the front of context + + +def _make_cross_attn(n_heads: int = 2, head_dim: int = 4) -> WanI2VCrossAttention: + """A WanI2VCrossAttention whose submodules are call-counting stubs. + + ``__init__`` is bypassed (it needs a TP group); we only wire what + ``forward`` touches. Each projection returns a [B, S, n*head_dim] tensor so + the subsequent ``unflatten(2, (n, head_dim))`` works, and is a MagicMock so + we can assert how many times it was invoked. + """ + inner = n_heads * head_dim + module = WanI2VCrossAttention.__new__(WanI2VCrossAttention) + module.tp_num_heads = n_heads + module.head_dim = head_dim + + def _proj() -> MagicMock: + return MagicMock(side_effect=lambda t: torch.zeros(t.shape[0], t.shape[1], inner)) + + module.q = _proj() + module.k = _proj() + module.v = _proj() + module.k_img = _proj() + module.v_img = _proj() + module.norm_q = lambda t: t + module.norm_k = lambda t: t + module.norm_k_img = lambda t: t + # attn returns the query unchanged ([B, S, n, d]); flatten(2) downstream is happy. + module.attn = lambda q, k, v: q + module.o = lambda t: t + return module + + +def _fresh_cache() -> dict: + return {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} + + +def _inputs(inner: int = 8): + x = torch.randn(1, 3, inner) # B=1, query_len=3 + context = torch.randn(1, _CLIP_LEN + 5, inner) # 257 image tokens + 5 text tokens + return x, context + + +def test_cache_inits_on_first_call_and_reuses_on_second() -> None: + module = _make_cross_attn() + cache = _fresh_cache() + x, context = _inputs() + + # First call: cache is initialised and every k/v projection runs once. + module.forward(x, context, crossattn_cache=cache) + assert cache["is_init"] is True + for key in ("k", "v", "k_img", "v_img"): + assert cache[key] is not None, f"{key} should be cached after first call" + assert module.k.call_count == 1 + assert module.v.call_count == 1 + assert module.k_img.call_count == 1 + assert module.v_img.call_count == 1 + + cached = {key: cache[key] for key in ("k", "v", "k_img", "v_img")} + + # Second call: k/v/k_img/v_img are read from cache, NOT recomputed. + module.forward(x, context, crossattn_cache=cache) + assert module.k.call_count == 1 + assert module.v.call_count == 1 + assert module.k_img.call_count == 1 + assert module.v_img.call_count == 1 + for key, tensor in cached.items(): + assert cache[key] is tensor, f"{key} should be the same cached object" + + # The query always depends on the per-step hidden state, so it is recomputed. + assert module.q.call_count == 2 + + +def test_no_cache_recomputes_every_call() -> None: + module = _make_cross_attn() + x, context = _inputs() + + module.forward(x, context, crossattn_cache=None) + module.forward(x, context, crossattn_cache=None) + + # Without a cache, the image/text k/v are recomputed on every call. + assert module.k_img.call_count == 2 + assert module.v_img.call_count == 2 + assert module.k.call_count == 2 + assert module.v.call_count == 2 diff --git a/tests/dreamzero/test_fused_qk_rms_norm.py b/tests/dreamzero/test_fused_qk_rms_norm.py new file mode 100644 index 000000000..dbbb43940 --- /dev/null +++ b/tests/dreamzero/test_fused_qk_rms_norm.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Equivalence tests for DreamZero's fused q/k qk-norm TP all-reduce. + +`fused_qk_rms_norm` packs the q and k per-token sum-of-squares into one tensor +and issues a single tensor-parallel all-reduce instead of two. These tests lock +in that it is numerically identical to applying `norm_q`/`norm_k` separately, +and that it really issues exactly one (fused) all-reduce. + +CPU-only: the TP world size and the collective are mocked, so a single process +simulates one rank (all-reduce is identity). This exercises the cat/split and +count-scaling logic that a real multi-rank run depends on. +""" + +from unittest.mock import patch + +import pytest +import torch + +from vllm_omni.diffusion.models.dreamzero.causal_wan_model import ( + DistributedRMSNorm, + fused_qk_rms_norm, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_MODULE = "vllm_omni.diffusion.models.dreamzero.causal_wan_model" + + +def _make_norms(dim: int) -> tuple[DistributedRMSNorm, DistributedRMSNorm]: + """Two norms with distinct random weights so a swapped q/k slice is caught.""" + torch.manual_seed(0) + norm_q = DistributedRMSNorm(dim) + norm_k = DistributedRMSNorm(dim) + norm_q.weight.data.normal_() + norm_k.weight.data.normal_() + return norm_q, norm_k + + +def _identity_all_reduce(tensor: torch.Tensor) -> torch.Tensor: + """Single-rank all-reduce: the reduction over one rank is the identity.""" + return tensor + + +def test_fused_matches_separate_tp1() -> None: + dim = 64 + norm_q, norm_k = _make_norms(dim) + q = torch.randn(2, 5, dim) + k = torch.randn(2, 5, dim) + + with patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1): + ref_q, ref_k = norm_q(q), norm_k(k) + fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k) + + torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0) + torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0) + + +def test_fused_matches_separate_tp_gt1() -> None: + dim = 64 + norm_q, norm_k = _make_norms(dim) + q = torch.randn(2, 5, dim) + k = torch.randn(2, 5, dim) + + # tp_size=2 exercises the reduce branch and the `global_count *= tp_size` + # path in both the per-norm forward and the fused helper. + with ( + patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2), + patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce), + ): + ref_q, ref_k = norm_q(q), norm_k(k) + fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k) + + torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0) + torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0) + + +def test_fused_issues_single_all_reduce_over_packed_pair() -> None: + dim = 64 + norm_q, norm_k = _make_norms(dim) + q = torch.randn(2, 5, dim) + k = torch.randn(2, 5, dim) + + with ( + patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2), + patch( + f"{_MODULE}.tensor_model_parallel_all_reduce", + side_effect=_identity_all_reduce, + ) as mock_ar, + ): + fused_qk_rms_norm(norm_q, norm_k, q, k) + + # Exactly one collective, over a tensor whose last dim packs the q+k pair. + assert mock_ar.call_count == 1 + packed = mock_ar.call_args[0][0] + assert packed.shape[-1] == 2 + + +def test_fused_identity_fallback_when_qk_norm_disabled() -> None: + # qk_norm=False makes the norms nn.Identity; the helper must fall back to + # plain passthrough without touching the collective. + norm_q, norm_k = torch.nn.Identity(), torch.nn.Identity() + q = torch.randn(2, 5, 64) + k = torch.randn(2, 5, 64) + + with patch(f"{_MODULE}.tensor_model_parallel_all_reduce") as mock_ar: + fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k) + + torch.testing.assert_close(fused_q, q, rtol=0, atol=0) + torch.testing.assert_close(fused_k, k, rtol=0, atol=0) + mock_ar.assert_not_called() diff --git a/tests/dreamzero/test_qkv_fusion.py b/tests/dreamzero/test_qkv_fusion.py new file mode 100644 index 000000000..16b78693f --- /dev/null +++ b/tests/dreamzero/test_qkv_fusion.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Equivalence test for DreamZero self-attn QKV fusion. + +CausalWanSelfAttention now projects q/k/v with a single QKVParallelLinear and +splits the result, instead of three separate ColumnParallelLinear. The DreamZero +checkpoint still ships separate q/k/v weights, so load_weights routes each into +the packed `qkv` param with its shard id. This test proves the chain +"load separate q/k/v via shard id -> single GEMM -> split" reproduces the three +independent projections, given identical weights. + +CPU-only: a single-rank (world_size=1) tensor-parallel group is initialized and +the GEMM dispatcher is forced to the default (CPU-compatible) path, following +tests/diffusion/layers/test_adalayernorm.py. +""" + +import os + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture(autouse=True) +def _init_distributed(): + """Minimal world_size=1 TP group required to build QKVParallelLinear.""" + from vllm.distributed.parallel_state import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, + ) + + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29507") + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method="env://", + ) + initialize_model_parallel() + yield + cleanup_dist_env_and_memory() + + +@pytest.fixture(autouse=True) +def _force_default_gemm(monkeypatch): + """Force CPU-compatible GEMM dispatch for CPU test tensors.""" + from vllm.model_executor.layers.utils import default_unquantized_gemm + + monkeypatch.setattr( + "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", + lambda: default_unquantized_gemm, + ) + + +def test_fused_qkv_load_and_split_matches_separate_projections(): + from vllm.model_executor.layers.linear import QKVParallelLinear + + torch.manual_seed(0) + dim, head_dim, num_heads = 16, 4, 4 # inner_dim == dim, no GQA + + qkv = QKVParallelLinear( + hidden_size=dim, + head_size=head_dim, + total_num_heads=num_heads, + bias=True, + ) + + # Distinct random weight + bias per projection, loaded the exact way + # DreamZero.load_weights does: param.weight_loader(param, tensor, shard_id). + refs: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + for shard_id in ("q", "k", "v"): + w = torch.randn(dim, dim) + b = torch.randn(dim) + refs[shard_id] = (w, b) + qkv.weight.weight_loader(qkv.weight, w, shard_id) + qkv.bias.weight_loader(qkv.bias, b, shard_id) + + x = torch.randn(2, 3, dim) + fused, _ = qkv(x) + q, k, v = fused.split([dim, dim, dim], dim=-1) + + for got, shard_id in zip((q, k, v), ("q", "k", "v"), strict=True): + w, b = refs[shard_id] + ref = torch.nn.functional.linear(x, w, b) + torch.testing.assert_close(got, ref) + + +def test_fused_qkv_shard_ids_are_not_transposed(): + # Guard against q/k/v being routed to the wrong slice: give each projection a + # constant, distinct weight so a swapped shard id would surface immediately. + from vllm.model_executor.layers.linear import QKVParallelLinear + + dim, head_dim, num_heads = 16, 4, 4 + qkv = QKVParallelLinear( + hidden_size=dim, + head_size=head_dim, + total_num_heads=num_heads, + bias=True, + ) + + scale = {"q": 1.0, "k": 2.0, "v": 3.0} + for shard_id, s in scale.items(): + qkv.weight.weight_loader(qkv.weight, torch.full((dim, dim), s), shard_id) + qkv.bias.weight_loader(qkv.bias, torch.zeros(dim), shard_id) + + x = torch.ones(1, dim) + fused, _ = qkv(x) + q, k, v = fused.split([dim, dim, dim], dim=-1) + + # F.linear(ones, full(s)) == s * dim per element. + torch.testing.assert_close(q, torch.full((1, dim), scale["q"] * dim)) + torch.testing.assert_close(k, torch.full((1, dim), scale["k"] * dim)) + torch.testing.assert_close(v, torch.full((1, dim), scale["v"] * dim)) diff --git a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py index e187054e2..ad3175e1e 100644 --- a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py +++ b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py @@ -21,11 +21,12 @@ import torch.nn as nn from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, - get_tp_group, + tensor_model_parallel_all_reduce, ) from vllm.model_executor.layers.conv import Conv3dLayer from vllm.model_executor.layers.linear import ( ColumnParallelLinear, + QKVParallelLinear, RowParallelLinear, ) from vllm.model_executor.utils import set_weight_attrs @@ -174,22 +175,82 @@ class DistributedRMSNorm(nn.Module): raise ValueError(f"RMSNorm shard shape mismatch: param={tuple(param.shape)}, shard={tuple(shard.shape)}.") param.data.copy_(shard) - def forward(self, x: torch.Tensor) -> torch.Tensor: - tp_size = get_tensor_model_parallel_world_size() + def _local_sum_sq(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]: + """Per-rank float32 activation, local sum-of-squares, and local width.""" x_float = x.float() local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True) - local_count = x.shape[-1] + return x_float, local_sum_sq, x.shape[-1] + + def _scale( + self, + x_float: torch.Tensor, + global_sum_sq: torch.Tensor, + global_count: int, + x: torch.Tensor, + ) -> torch.Tensor: + """Apply the (already reduced) RMS and the per-rank weight shard.""" + mean_sq = global_sum_sq / global_count + return (x_float * torch.rsqrt(mean_sq + self.eps)).type_as(x) * self.weight + + def forward(self, x: torch.Tensor) -> torch.Tensor: + tp_size = get_tensor_model_parallel_world_size() + x_float, local_sum_sq, local_count = self._local_sum_sq(x) if tp_size > 1: - global_sum_sq = local_sum_sq.clone() - torch.distributed.all_reduce(global_sum_sq, group=get_tp_group().device_group) + # Use vLLM's collective (custom all-reduce / symmetric-mem fast path + # for small tensors) instead of raw torch.distributed.all_reduce, and + # take the return value so the custom-AR path (which may return a new + # buffer) is handled correctly. No .clone() needed. + global_sum_sq = tensor_model_parallel_all_reduce(local_sum_sq) global_count = local_count * tp_size else: global_sum_sq = local_sum_sq global_count = local_count - mean_sq = global_sum_sq / global_count - return (x_float * torch.rsqrt(mean_sq + self.eps)).type_as(x) * self.weight + return self._scale(x_float, global_sum_sq, global_count, x) + + +def fused_qk_rms_norm( + norm_q: nn.Module, + norm_k: nn.Module, + q: torch.Tensor, + k: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply q/k :class:`DistributedRMSNorm` with a SINGLE fused TP all-reduce. + + WHY: self-attention norms q and k every step; on TP each norm issues its own + tiny all-reduce of the per-token sum-of-squares. For DreamZero's decode-like + per-step forward these latency-bound micro-collectives dominate, so we pack + both sum-of-squares into one tensor and reduce once (2 collectives → 1). + + NUMERICALLY IDENTICAL to ``norm_q(q), norm_k(k)``: all-reduce is elementwise, + so packing along the last dim reduces each slice independently with the same + fp32 accumulation. Requires q and k to share the same shape (true for + self-attention — both come from the same hidden states). Falls back to + independent application when either norm is not a DistributedRMSNorm + (e.g. nn.Identity when qk_norm=False). + """ + if not (isinstance(norm_q, DistributedRMSNorm) and isinstance(norm_k, DistributedRMSNorm)): + return norm_q(q), norm_k(k) + + # The fused path reduces one packed sum-of-squares and reuses q's token width + # as the RMS count for BOTH q and k, so q and k must share the same shape. + # (Self-attention guarantees this: q and k are projected from the same x.) + assert q.shape == k.shape, "fused_qk_rms_norm requires q and k to have the same shape." + + tp_size = get_tensor_model_parallel_world_size() + q_float, q_sum_sq, count = norm_q._local_sum_sq(q) + k_float, k_sum_sq, _ = norm_k._local_sum_sq(k) + + if tp_size > 1: + packed = torch.cat([q_sum_sq, k_sum_sq], dim=-1) + packed = tensor_model_parallel_all_reduce(packed) + q_sum_sq, k_sum_sq = packed[..., 0:1], packed[..., 1:2] + count = count * tp_size + + q_out = norm_q._scale(q_float, q_sum_sq, count, q) + k_out = norm_k._scale(k_float, k_sum_sq, count, k) + return q_out, k_out # ── Projections ───────────────────────────────────────────────────── @@ -337,22 +398,35 @@ class WanI2VCrossAttention(nn.Module): context = context[:, 257:] n, d = self.tp_num_heads, self.head_dim q = self.norm_q(self.q(x)).unflatten(2, (n, d)) + # context (text) and context_img (clip features) are constant within a + # session, so k/v and k_img/v_img are cached on first call and reused. + # context_img == img_emb(clip feature) is session-invariant (state.clip_feas + # is set once at current_start_frame==0 and only cleared on reset, which also + # rebuilds crossattn_cache). Caching k_img/v_img removes a per-step + # norm_k_img all-reduce plus the k_img/v_img projection GEMMs in steady state. + # Only q depends on the per-step x and is always recomputed. if crossattn_cache is not None: if not crossattn_cache["is_init"]: crossattn_cache["is_init"] = True k = self.norm_k(self.k(context)).unflatten(2, (n, d)) v = self.v(context).unflatten(2, (n, d)) + k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d)) + v_img = self.v_img(context_img).unflatten(2, (n, d)) crossattn_cache["k"] = k crossattn_cache["v"] = v + crossattn_cache["k_img"] = k_img + crossattn_cache["v_img"] = v_img else: k = crossattn_cache["k"] v = crossattn_cache["v"] + k_img = crossattn_cache["k_img"] + v_img = crossattn_cache["v_img"] else: k = self.norm_k(self.k(context)).unflatten(2, (n, d)) v = self.v(context).unflatten(2, (n, d)) + k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d)) + v_img = self.v_img(context_img).unflatten(2, (n, d)) x = self.attn(q, k, v) - k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d)) - v_img = self.v_img(context_img).unflatten(2, (n, d)) img_x = self.attn(q, k_img, v_img) x = x.flatten(2) img_x = img_x.flatten(2) @@ -403,9 +477,21 @@ class CausalWanSelfAttention(nn.Module): self.num_action_per_block = num_action_per_block self.num_state_per_block = num_state_per_block self.max_attention_size = 21 * frame_seqlen if local_attn_size == -1 else local_attn_size * frame_seqlen - self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False) - self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False) - self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False) + # Fused QKV projection: q/k/v all come from x, so a single column-parallel + # GEMM replaces three (fewer launches / better GEMM util in the decode-like + # per-step forward). No GQA here: total_num_kv_heads defaults to num_heads. + self.qkv = QKVParallelLinear( + hidden_size=dim, + head_size=self.head_dim, + total_num_heads=num_heads, + bias=True, + ) + # The forward splits qkv into three equal q/k/v shards, which is only + # correct without GQA. total_num_kv_heads defaults to total_num_heads + # here; fail loud if a future checkpoint introduces GQA so the equal + # split does not silently misalign k/v. + if self.qkv.total_num_kv_heads != self.qkv.total_num_heads: + raise ValueError("Self-attn QKV fusion requires no GQA (total_num_kv_heads == total_num_heads).") self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True, return_bias=False) self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity() self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity() @@ -431,9 +517,15 @@ class CausalWanSelfAttention(nn.Module): """Inference-only forward (KV cache path).""" n, d = self.tp_num_heads, self.head_dim - q = self.norm_q(self.q(x)).unflatten(2, (n, d)) - k = self.norm_k(self.k(x)).unflatten(2, (n, d)) - v = self.v(x).unflatten(2, (n, d)) + # Single fused QKV GEMM, then split into the per-rank q/k/v shards. + qkv, _ = self.qkv(x) + qk_size = self.tp_num_heads * self.head_dim + q, k, v = qkv.split([qk_size, qk_size, qk_size], dim=-1) + # Fuse q/k qk-norm into a single TP all-reduce (both come from x, same shape). + q, k = fused_qk_rms_norm(self.norm_q, self.norm_k, q, k) + q = q.unflatten(2, (n, d)) + k = k.unflatten(2, (n, d)) + v = v.unflatten(2, (n, d)) updated_kv_cache: torch.Tensor | None = None @@ -834,6 +926,7 @@ class CausalWanModel(nn.Module): timestep_action: torch.Tensor | None, state: torch.Tensor | None, kv_cache: list[torch.Tensor], + crossattn_cache: list[dict] | None, current_start_frame: int, ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]: x = x.flatten(start_dim=2).transpose(1, 2) @@ -884,6 +977,7 @@ class CausalWanModel(nn.Module): context=context, action_register_length=action_register_length, kv_cache=kv_cache[block_index] if kv_cache else None, + crossattn_cache=crossattn_cache[block_index] if crossattn_cache else None, current_start_frame=current_start_frame, ) updated_kv_caches.append(updated_kv_cache) @@ -941,6 +1035,7 @@ class CausalWanModel(nn.Module): timestep_action=timestep_action, state=state, kv_cache=kv_cache, + crossattn_cache=crossattn_cache, current_start_frame=current_start_frame, ) diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py index 30ea2a4ce..1d1429534 100644 --- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py +++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py @@ -1104,10 +1104,27 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): .replace("img_emb.proj.3.", "img_emb.fc2.") .replace("img_emb.proj.4.", "img_emb.norm2.") ) + + # Self-attn q/k/v are fused into a single QKVParallelLinear; route + # each separate checkpoint weight/bias to the packed `qkv` param + # with its shard id. cross_attn keeps separate q/k/v (q from x, + # k/v from context — not fusible), so it is left untouched here. + qkv_shard_id: str | None = None + for shard_id in ("q", "k", "v"): + needle = f".self_attn.{shard_id}." + if needle in new_name: + new_name = new_name.replace(needle, ".self_attn.qkv.") + qkv_shard_id = shard_id + break + if new_name in params: param = params[new_name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, tensor) + if qkv_shard_id is not None: + # QKVParallelLinear.weight_loader needs the shard id. + param.weight_loader(param, tensor, qkv_shard_id) + else: + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, tensor) loaded.add(new_name) elif new_name in buffers: buffers[new_name].data.copy_(tensor) diff --git a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py index 0972d1b5a..9c1e7873d 100644 --- a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py +++ b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py @@ -131,8 +131,12 @@ class DreamZeroState: torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers) ] - self.crossattn_cache = [{"is_init": False, "k": None, "v": None} for _ in range(num_layers)] - self.crossattn_cache_neg = [{"is_init": False, "k": None, "v": None} for _ in range(num_layers)] + self.crossattn_cache = [ + {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers) + ] + self.crossattn_cache_neg = [ + {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers) + ] def update_kv_cache( self, From 1228a109f09516c4e9323e9d1774b76304a84047 Mon Sep 17 00:00:00 2001 From: yangyonggit Date: Thu, 18 Jun 2026 20:11:51 +0200 Subject: [PATCH 09/40] [Bugfix] MOSS-TTS: fix cross-request audio corruption under batching (#4415) Signed-off-by: leo.yang Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> --- .../models/moss_tts/test_per_request_codes.py | 176 ++++++++++++++++++ .../moss_tts/modeling_moss_tts_talker.py | 79 +++++--- vllm_omni/worker/gpu_model_runner.py | 14 ++ 3 files changed, 245 insertions(+), 24 deletions(-) create mode 100644 tests/model_executor/models/moss_tts/test_per_request_codes.py diff --git a/tests/model_executor/models/moss_tts/test_per_request_codes.py b/tests/model_executor/models/moss_tts/test_per_request_codes.py new file mode 100644 index 000000000..eabb4fdc0 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_per_request_codes.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for MOSS-TTS per-request code routing in the talker. + +Guards the cross-request audio-corruption class (RFC #4316 Issue 1, bug report +#4355): under continuous batching every Stage-1 request must decode its *own* +codes, not request 0's. The talker emits ``codes["audio"]`` as a batch-ordered +**list** so the generic splitter (``to_payload_element``) routes +``element[idx]`` to request ``idx``; the list must stay index-aligned with the +batch (placeholder for skipped requests) or a shorter list silently falls back +to entry 0 — the original corruption. The hidden-row mapping must also use the +runner's real per-request token spans, not an equal split, or mixed +prefill+decode steps sample each request's codes from the wrong rows. + +CPU-only and weight-free: ``make_omni_output`` is driven directly on a bare +talker with ``_sample_audio_codes`` mocked to echo the hidden row it was given, +so an assertion on the output value proves which row each request sampled from. +""" + +from __future__ import annotations + +import functools + +import pytest + +torch = pytest.importorskip("torch") + +N_VQ = 4 + + +@functools.lru_cache(maxsize=1) +def _delay_talker_cls(): + """Defer talker import (pulls vLLM model_executor) until first use.""" + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import ( + MossTTSDelayTalkerForGeneration, + ) + + return MossTTSDelayTalkerForGeneration + + +def _make_bare_delay_talker(): + """A talker whose only behaviour is the routing/span logic under test. + + ``_sample_audio_codes`` is replaced with a stub returning an (N_VQ,) tensor + filled with the value of the hidden row it received. With a ``hidden`` where + row r holds the scalar r, the emitted codes reveal which row each request + sampled. + """ + cls = _delay_talker_cls() + talker = cls.__new__(cls) + talker._batch_state = None + talker.n_vq = N_VQ + + def _stub_sample(last_h, state): + row_value = int(last_h.reshape(-1)[0].item()) + return torch.full((N_VQ,), row_value, dtype=torch.long) + + talker._sample_audio_codes = _stub_sample + return talker + + +def _row_indexed_hidden(num_rows: int) -> torch.Tensor: + """(num_rows, 1) tensor where row r holds the scalar r.""" + return torch.arange(num_rows, dtype=torch.float32).reshape(num_rows, 1) + + +def _audio_payload(result): + return result.multimodal_outputs["codes"]["audio"] + + +class TestPerRequestRouting: + def test_mixed_prefill_decode_routes_by_span(self): + """req 0 prefills 3 rows, req 1 decodes 1 row in the same step. + + With real spans, req 0 samples its last row (2) and req 1 its last row + (3); the payload is a batch-aligned list, element[idx] per request. + """ + talker = _make_bare_delay_talker() + hidden = _row_indexed_hidden(4) + info_dicts = [{}, {}] + spans = [(0, 3), (3, 4)] + + result = talker.make_omni_output( + hidden, + runtime_additional_information=info_dicts, + request_token_spans=spans, + ) + audio = _audio_payload(result) + + assert isinstance(audio, list) + assert len(audio) == 2 # one entry per request, batch-aligned + assert audio[0].shape == (1, N_VQ) + assert torch.all(audio[0] == 2) # req 0: last row of span [0,3) + assert torch.all(audio[1] == 3) # req 1: last row of span [3,4) + # No cross-request leakage. + assert not torch.equal(audio[0], audio[1]) + + def test_no_request_samples_outside_its_span(self): + talker = _make_bare_delay_talker() + spans = [(0, 2), (2, 5), (5, 6)] + hidden = _row_indexed_hidden(6) + info_dicts = [{}, {}, {}] + + result = talker.make_omni_output( + hidden, + runtime_additional_information=info_dicts, + request_token_spans=spans, + ) + audio = _audio_payload(result) + + assert len(audio) == 3 + sampled = [] + for i, (start, end) in enumerate(spans): + value = int(audio[i].reshape(-1)[0].item()) + assert start <= value < end, f"req {i} sampled row {value} outside [{start},{end})" + sampled.append(value) + assert len(set(sampled)) == len(sampled) # distinct rows, no sharing + + def test_skipped_request_keeps_list_aligned(self): + """An empty-span request gets an empty placeholder at its own index, so + later requests are not shifted onto request 0's codes.""" + talker = _make_bare_delay_talker() + # req 1 has an empty span (start == end) and must not displace req 2. + spans = [(0, 1), (1, 1), (1, 2)] + hidden = _row_indexed_hidden(2) + info_dicts = [{}, {}, {}] + + result = talker.make_omni_output( + hidden, + runtime_additional_information=info_dicts, + request_token_spans=spans, + ) + audio = _audio_payload(result) + + assert len(audio) == 3 + assert torch.all(audio[0] == 0) # req 0: row 0 + assert audio[1].numel() == 0 # req 1: empty placeholder, aligned + assert torch.all(audio[2] == 1) # req 2: row 1, not shifted + + def test_accumulates_into_per_request_history(self): + talker = _make_bare_delay_talker() + hidden = _row_indexed_hidden(2) + prior = torch.full((1, N_VQ), 7, dtype=torch.long) + info_dicts = [{"audio_codes": {"accumulated": prior}}, {}] + spans = [(0, 1), (1, 2)] + + result = talker.make_omni_output( + hidden, + runtime_additional_information=info_dicts, + request_token_spans=spans, + ) + audio = _audio_payload(result) + + assert audio[0].shape == (2, N_VQ) + assert torch.all(audio[0][0] == 7) # prior history kept + assert torch.all(audio[0][1] == 0) # this step's row 0 appended + assert torch.all(audio[1] == 1) + + def test_runs_without_spans_backward_compat(self): + """Missing request_token_spans must not crash (equal-split fallback) + and still yields a batch-aligned list.""" + talker = _make_bare_delay_talker() + hidden = _row_indexed_hidden(2) + info_dicts = [{}, {}] + + result = talker.make_omni_output(hidden, runtime_additional_information=info_dicts) + audio = _audio_payload(result) + assert isinstance(audio, list) + assert len(audio) == 2 + + def test_no_codes_returns_empty_payload(self): + """Zero hidden rows → no codes → empty multimodal payload (no list).""" + talker = _make_bare_delay_talker() + hidden = torch.zeros((0, 1), dtype=torch.float32) + result = talker.make_omni_output(hidden, runtime_additional_information=[{}], request_token_spans=[(0, 0)]) + assert result.multimodal_outputs == {} diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py index e7e100a18..a03b4c1d1 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py @@ -553,7 +553,19 @@ class MossTTSDelayTalkerForGeneration(nn.Module): # rows align with hidden rows, which align with info_dicts in order. self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts] - audio_codes_list: list[torch.Tensor] = [] + # Per-request (start, end) hidden-row spans from the runner. In mixed + # prefill+decode steps the per-request token counts differ, so an equal + # rows-per-request split would sample codes from the wrong rows; fall + # back to that split only when spans are unavailable. + spans = kwargs.get("request_token_spans") + + # One accumulated-codes tensor per request, in batch order. The generic + # splitter (``to_payload_element``) routes ``per_req_codes[idx]`` to + # request ``idx``, so Stage 1 decodes each request's own codes instead + # of request 0's. Skipped requests get an empty placeholder so indices + # stay aligned (a shorter list would silently fall back to entry 0). + per_req_codes: list[torch.Tensor] = [] + have_codes = False if hidden.numel() > 0: num_rows = hidden.shape[0] @@ -561,10 +573,16 @@ class MossTTSDelayTalkerForGeneration(nn.Module): for i, info in enumerate(info_dicts): if not isinstance(info, dict): + per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long)) continue - row_start = i * rows_per_req - row_end = min(row_start + rows_per_req, num_rows) + if spans is not None and i < len(spans): + row_start, row_end = spans[i] + row_end = min(int(row_end), num_rows) + else: + row_start = i * rows_per_req + row_end = min(row_start + rows_per_req, num_rows) if row_start >= row_end: + per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long)) continue # Sample new audio codes from the last hidden state of this request. @@ -582,21 +600,20 @@ class MossTTSDelayTalkerForGeneration(nn.Module): "current": new_codes, "accumulated": updated_acc, } - audio_codes_list.append(updated_acc) + per_req_codes.append(updated_acc) + have_codes = True - if not audio_codes_list: + if not have_codes: return OmniOutput( text_hidden_states=hidden, multimodal_outputs={}, ) - # Forward the accumulated audio codes (T, NQ) of the first request to - # Stage 1. Multi-request batching for the codec is not yet supported - # by this adapter; the existing code already takes the first stage - # output anyway. + # Emit one accumulated-codes tensor per request as a list; the splitter + # picks ``element[idx]`` per request on the way to Stage 1. return OmniOutput( text_hidden_states=hidden, - multimodal_outputs={"codes": {"audio": audio_codes_list[0]}}, + multimodal_outputs={"codes": {"audio": per_req_codes}}, ) # ------------------------------------------------------------------ @@ -933,15 +950,27 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts] - audio_codes_list: list[torch.Tensor] = [] + # See delay talker: real per-request row spans from the runner, with the + # equal-split as fallback when unavailable. + spans = kwargs.get("request_token_spans") + # Per-request accumulated codes in batch order, pre-filled with empty + # placeholders so every skip path (stopped / eos / bos / pad) keeps + # indices aligned for to_payload_element's per-request routing. A + # shorter list would silently fall back to request 0's codes. + per_req_codes: list[torch.Tensor] = [hidden.new_empty((0, self.n_vq), dtype=torch.long) for _ in info_dicts] + have_codes = False if hidden.numel() > 0 and info_dicts: num_rows = hidden.shape[0] rows_per_req = max(1, num_rows // max(1, len(info_dicts) or 1)) for i, info in enumerate(info_dicts): if not isinstance(info, dict): continue - row_start = i * rows_per_req - row_end = min(row_start + rows_per_req, num_rows) + if spans is not None and i < len(spans): + row_start, row_end = spans[i] + row_end = min(int(row_end), num_rows) + else: + row_start = i * rows_per_req + row_end = min(row_start + rows_per_req, num_rows) if row_start >= row_end: continue @@ -1008,23 +1037,25 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc} state["step"] = int(state.get("step", 0)) + 1 - audio_codes_list.append(updated_acc) + per_req_codes[i] = updated_acc + have_codes = True - if not audio_codes_list: + if not have_codes: return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) - # Forward the first request's accumulated codes (multi-request - # batching for the codec is not yet wired up — same as delay). - # The realtime variant emits raw codes (no delay pattern), so we - # signal that to the chunk processor via a 1-D bool tensor in - # ``meta.finished`` (already a tensor field on the schema). The - # processor reads its truth value as "skip de-delay" only for - # realtime; the delay path doesn't populate ``meta`` at all. - device = audio_codes_list[0].device + # Emit one accumulated-codes tensor per request as a list; the splitter + # routes ``element[idx]`` to request idx, so each Stage-1 request decodes + # its own codes instead of request 0's. + # The realtime variant emits raw codes (no delay pattern), so we signal + # that to the chunk processor via a 1-D bool tensor in ``meta.finished`` + # (already a tensor field on the schema). The processor reads its truth + # value as "skip de-delay" only for realtime; the delay path doesn't + # populate ``meta`` at all. + device = hidden.device return OmniOutput( text_hidden_states=hidden, multimodal_outputs={ - "codes": {"audio": audio_codes_list[0]}, + "codes": {"audio": per_req_codes}, "meta": {"finished": torch.tensor([True], dtype=torch.bool, device=device)}, }, ) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 349e1da04..cdf2286bf 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1313,6 +1313,20 @@ class OmniGPUModelRunner(GPUModelRunner): traceback.print_exc() + # Per-request (start, end) hidden-row spans so make_omni_output can map + # flat hidden rows to the right request in mixed prefill+decode steps, + # instead of assuming an equal rows-per-request split (which samples the + # wrong rows whenever per-request token counts differ). + nstp = self._omni_num_scheduled_tokens_np + if nstp is not None and len(nstp) == len(self.input_batch.req_ids): + try: + model_kwargs_extra["request_token_spans"] = self._compute_request_token_spans(nstp) + except Exception as e: + # Visible on purpose: the fallback is the equal rows-per-request + # split, which can re-introduce the cross-request corruption this + # plumbing fixes — a silent failure here must not pass unnoticed. + logger.warning("[OMNI] Failed to compute request_token_spans: %s", e) + if self._omni_query_start_loc_model_kwarg: try: num_reqs = len(self.input_batch.req_ids) From 13e9094be092471036a2ef53e20b3b078bdaf60c Mon Sep 17 00:00:00 2001 From: yangyonggit Date: Fri, 19 Jun 2026 00:27:26 +0200 Subject: [PATCH 10/40] [Doc] Add 1x RTX A6000 48GB section to Voxtral TTS recipe (#2645) (#4051) Signed-off-by: leo.yang --- recipes/mistralai/Voxtral-TTS.md | 42 ++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/recipes/mistralai/Voxtral-TTS.md b/recipes/mistralai/Voxtral-TTS.md index 595d5c29e..aae48a689 100644 --- a/recipes/mistralai/Voxtral-TTS.md +++ b/recipes/mistralai/Voxtral-TTS.md @@ -105,3 +105,45 @@ voice-cloning e2e test ([`Voice cloning (capability gated upstream)`](../../exam ) or adding a `ref_audio` request parameter will fail because the public checkpoint does not provide the encoder weights needed to convert reference audio into conditioning features. The expected failure is: `RuntimeError: encode_waveforms requires encoder weights which are not available in the open-source checkpoint.` + +### 1x NVIDIA RTX A6000 48GB + +#### Environment + +- OS: Ubuntu 22.04 +- Python: 3.12 +- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0 +- GPU: NVIDIA RTX A6000, 48 GB +- vLLM version: 0.22.0 +- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d +- Required package: `pip install mistral-common` + +#### Command + +```bash +vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091 +``` + +#### Verification + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "mistralai/Voxtral-4B-TTS-2603", + "input": "Artificial intelligence is transforming the way we interact with technology.", + "voice": "neutral_female", + "response_format": "wav" + }' -o voxtral_test.wav +``` + +Supported voice presets: `ar_male`, `casual_female`, `casual_male`, `cheerful_female`, +`de_female`, `de_male`, `es_female`, `es_male`, `fr_female`, `fr_male`, `hi_female`, +`hi_male`, `it_female`, `it_male`, `neutral_female`, `neutral_male`, `nl_female`, +`nl_male`, `pt_female`, `pt_male`. + +#### Notes + +- Memory usage: Peak VRAM ~40.5 GiB during inference. +- Generation time: ~8 seconds for ~200 words of English text. +- Audio quality is high; mono 24 kHz output. From 37a442447ebb474451aa56abd6115d780a38d76e Mon Sep 17 00:00:00 2001 From: yangyonggit Date: Fri, 19 Jun 2026 00:32:18 +0200 Subject: [PATCH 11/40] [Cleanup] Remove duplicate HiDreamI1ImagePipeline (fixes #4009) (#4045) Signed-off-by: leo.yang --- .../pipeline_hidream_i1_image.py | 1059 ----------------- 1 file changed, 1059 deletions(-) delete mode 100644 vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py diff --git a/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py b/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py deleted file mode 100644 index abc50eb8a..000000000 --- a/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py +++ /dev/null @@ -1,1059 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import inspect -import json -import math -import os -from collections.abc import Callable, Iterable -from typing import Any - -import numpy as np -import torch -from diffusers.image_processor import VaeImageProcessor -from diffusers.models import AutoencoderKL -from diffusers.schedulers import FlowMatchEulerDiscreteScheduler -from diffusers.utils import deprecate, logging -from diffusers.utils.torch_utils import randn_tensor -from torch import nn -from transformers import ( - CLIPTextModelWithProjection, - CLIPTokenizer, - LlamaForCausalLM, - PreTrainedTokenizerFast, - T5EncoderModel, - T5Tokenizer, -) -from vllm.model_executor.models.utils import AutoWeightsLoader - -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin -from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders -from vllm_omni.diffusion.models.hidream_image import HiDreamImageTransformer2DModel -from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin -from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs -from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific - -logger = logging.get_logger(__name__) - - -def get_hidream_image_post_process_func( - od_config: OmniDiffusionConfig, -): - model_name = od_config.model - if os.path.exists(model_name): - model_path = model_name - else: - model_path = download_weights_from_hf_specific(model_name, None, ["*"]) - vae_config_path = os.path.join(model_path, "vae/config.json") - if not os.path.exists(vae_config_path): - raise FileNotFoundError( - f"VAE config not found at {vae_config_path}. " - "Please ensure the model path contains a valid VAE configuration." - ) - with open(vae_config_path) as f: - vae_config = json.load(f) - vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8 - - image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2) - - def post_process_func( - images: torch.Tensor, - ): - return image_processor.postprocess(images) - - return post_process_func - - -# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift -def calculate_shift( - image_seq_len, - base_seq_len: int = 256, - max_seq_len: int = 4096, - base_shift: float = 0.5, - max_shift: float = 1.15, -): - m = (max_shift - base_shift) / (max_seq_len - base_seq_len) - b = base_shift - m * base_seq_len - mu = image_seq_len * m + b - return mu - - -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps -def retrieve_timesteps( - scheduler, - num_inference_steps: int | None = None, - device: str | torch.device | None = None, - timesteps: list[int] | None = None, - sigmas: list[float] | None = None, - **kwargs, -): - r""" - Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles - custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. - - Args: - scheduler (`SchedulerMixin`): - The scheduler to get timesteps from. - num_inference_steps (`int`): - The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` - must be `None`. - device (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`list[int]`, *optional*): - Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, - `num_inference_steps` and `sigmas` must be `None`. - sigmas (`list[float]`, *optional*): - Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, - `num_inference_steps` and `timesteps` must be `None`. - - Returns: - timesteps (`torch.Tensor`): The timestep schedule from the scheduler. - num_inference_steps (`int`): The number of inference steps. - """ - if timesteps is not None and sigmas is not None: - raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") - if timesteps is not None: - accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) - if not accepts_timesteps: - raise ValueError( - f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" - f" timestep schedules. Please check whether you are using the correct scheduler." - ) - scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) - timesteps = scheduler.timesteps - num_inference_steps = len(timesteps) - elif sigmas is not None: - accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) - if not accept_sigmas: - raise ValueError( - f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" - f" sigmas schedules. Please check whether you are using the correct scheduler." - ) - scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) - timesteps = scheduler.timesteps - num_inference_steps = len(timesteps) - else: - scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) - timesteps = scheduler.timesteps - return timesteps, num_inference_steps - - -class HiDreamI1ImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, ProgressBarMixin): - def __init__( - self, - *, - od_config: OmniDiffusionConfig, - prefix: str = "", - ): - super().__init__() - self.od_config = od_config - dtype = getattr(od_config, "dtype", torch.bfloat16) - llama_path = self.od_config.extras["auxiliary_text_encoder"] - if llama_path is None: - logger.warning( - f"auxiliary_text_encoder is not provided. " - f"Attempting to load default LLaMA model: {self.od_config.extras['default_llama_model_id']}" - ) - llama_path = self.od_config.extras["default_llama_model_id"] - logger.info(f"Loading LLAMA model from: {llama_path}") - - self.weights_sources = [ - DiffusersPipelineLoader.ComponentSource( - model_or_path=od_config.model, - subfolder="transformer", - revision=None, - prefix="transformer.", - fall_back_to_pt=True, - ), - ] - - self.device = get_local_device() - model = od_config.model - # Check if model is a local path - local_files_only = os.path.exists(model) - - # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder - # race; prefetch the in-repo component set before any from_pretrained - # (``text_encoder_4`` lives in a separate Llama repo and is unaffected). - hidream_subfolders = [ - "scheduler", - "vae", - "text_encoder", - "tokenizer", - "text_encoder_2", - "tokenizer_2", - "text_encoder_3", - "tokenizer_3", - ] - prefetch_subfolders(model, hidream_subfolders, local_files_only=local_files_only) - - self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( - model, subfolder="scheduler", local_files_only=local_files_only - ) - self.vae = from_pretrained_with_prefetch( - AutoencoderKL.from_pretrained, - model, - subfolder="vae", - prefetch_list=hidream_subfolders, - local_files_only=local_files_only, - ).to(self.device) - self.text_encoder = from_pretrained_with_prefetch( - CLIPTextModelWithProjection.from_pretrained, - model, - subfolder="text_encoder", - prefetch_list=hidream_subfolders, - local_files_only=local_files_only, - ) - self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) - self.text_encoder_2 = from_pretrained_with_prefetch( - CLIPTextModelWithProjection.from_pretrained, - model, - subfolder="text_encoder_2", - prefetch_list=hidream_subfolders, - local_files_only=local_files_only, - ) - self.tokenizer_2 = CLIPTokenizer.from_pretrained( - model, subfolder="tokenizer_2", local_files_only=local_files_only - ) - self.text_encoder_3 = from_pretrained_with_prefetch( - T5EncoderModel.from_pretrained, - model, - subfolder="text_encoder_3", - prefetch_list=hidream_subfolders, - local_files_only=local_files_only, - ) - self.tokenizer_3 = T5Tokenizer.from_pretrained( - model, subfolder="tokenizer_3", local_files_only=local_files_only - ) - self.text_encoder_4 = LlamaForCausalLM.from_pretrained(llama_path, output_hidden_states=True, dtype=dtype).to( - self.device - ) - self.tokenizer_4 = PreTrainedTokenizerFast.from_pretrained(llama_path, use_fast=False) - transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HiDreamImageTransformer2DModel) - self.transformer = HiDreamImageTransformer2DModel( - od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs - ) - - self.stage = None - - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - # HiDreamImage latents are turned into 2x2 patches and packed. - # This means the latent width and height has to be divisible - # by the patch size. So the vae scale factor is multiplied by the patch size to account for this - self.default_sample_size = 128 - if getattr(self, "tokenizer_4", None) is not None: - self.tokenizer_4.pad_token = self.tokenizer_4.eos_token - - self.setup_diffusion_pipeline_profiler( - enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler - ) - - def check_inputs( - self, - prompt, - prompt_2, - prompt_3, - prompt_4, - negative_prompt=None, - negative_prompt_2=None, - negative_prompt_3=None, - negative_prompt_4=None, - prompt_embeds_t5=None, - prompt_embeds_llama3=None, - negative_prompt_embeds_t5=None, - negative_prompt_embeds_llama3=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - ): - if prompt is not None and pooled_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `pooled_prompt_embeds`: {pooled_prompt_embeds}. " - "Please make sure to only forward one of the two." - ) - elif prompt_2 is not None and pooled_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt_2`: {prompt_2} and `pooled_prompt_embeds`: {pooled_prompt_embeds}." - "Please make sure to only forward one of the two." - ) - elif prompt_3 is not None and prompt_embeds_t5 is not None: - raise ValueError( - f"Cannot forward both `prompt_3`: {prompt_3} and `prompt_embeds_t5`: {prompt_embeds_t5}." - "Please make sure to only forward one of the two." - ) - elif prompt_4 is not None and prompt_embeds_llama3 is not None: - raise ValueError( - f"Cannot forward both `prompt_4`: {prompt_4} and `prompt_embeds_llama3`: {prompt_embeds_llama3}." - "Please make sure to only forward one of the two." - ) - elif prompt is None and pooled_prompt_embeds is None: - raise ValueError( - "Provide either `prompt` or `pooled_prompt_embeds`." - "Cannot leave both `prompt` and `pooled_prompt_embeds` undefined." - ) - elif prompt is None and prompt_embeds_t5 is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds_t5`." - "Cannot leave both `prompt` and `prompt_embeds_t5` undefined." - ) - elif prompt is None and prompt_embeds_llama3 is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds_llama3`." - "Cannot leave both `prompt` and `prompt_embeds_llama3` undefined." - ) - elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)): - raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}") - elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)): - raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}") - elif prompt_4 is not None and (not isinstance(prompt_4, str) and not isinstance(prompt_4, list)): - raise ValueError(f"`prompt_4` has to be of type `str` or `list` but is {type(prompt_4)}") - - if negative_prompt is not None and negative_pooled_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_pooled_prompt_embeds`:" - f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two." - ) - elif negative_prompt_2 is not None and negative_pooled_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_pooled_prompt_embeds`:" - f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two." - ) - elif negative_prompt_3 is not None and negative_prompt_embeds_t5 is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} and `negative_prompt_embeds_t5`:" - f" {negative_prompt_embeds_t5}. Please make sure to only forward one of the two." - ) - elif negative_prompt_4 is not None and negative_prompt_embeds_llama3 is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_4`: {negative_prompt_4} and `negative_prompt_embeds_llama3`:" - f" {negative_prompt_embeds_llama3}. Please make sure to only forward one of the two." - ) - - if pooled_prompt_embeds is not None and negative_pooled_prompt_embeds is not None: - if pooled_prompt_embeds.shape != negative_pooled_prompt_embeds.shape: - raise ValueError( - "`pooled_prompt_embeds` and `negative_pooled_prompt_embeds`" - " must have the same shape when passed directly, but" - f" got: `pooled_prompt_embeds` {pooled_prompt_embeds.shape} != `negative_pooled_prompt_embeds`" - f" {negative_pooled_prompt_embeds.shape}." - ) - if prompt_embeds_t5 is not None and negative_prompt_embeds_t5 is not None: - if prompt_embeds_t5.shape != negative_prompt_embeds_t5.shape: - raise ValueError( - "`prompt_embeds_t5` and `negative_prompt_embeds_t5`" - " must have the same shape when passed directly, but" - f" got: `prompt_embeds_t5` {prompt_embeds_t5.shape} != `negative_prompt_embeds_t5`" - f" {negative_prompt_embeds_t5.shape}." - ) - if prompt_embeds_llama3 is not None and negative_prompt_embeds_llama3 is not None: - if prompt_embeds_llama3.shape != negative_prompt_embeds_llama3.shape: - raise ValueError( - "`prompt_embeds_llama3` and `negative_prompt_embeds_llama3`" - " must have the same shape when passed directly, but" - f" got: `prompt_embeds_llama3` {prompt_embeds_llama3.shape} != `negative_prompt_embeds_llama3`" - f" {negative_prompt_embeds_llama3.shape}." - ) - - def _get_t5_prompt_embeds( - self, - prompt: str | list[str] = None, - max_sequence_length: int = 128, - dtype: torch.dtype | None = None, - ): - dtype = dtype or self.text_encoder_3.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - - text_inputs = self.tokenizer_3( - prompt, - padding="max_length", - max_length=min(max_sequence_length, self.tokenizer_3.model_max_length), - truncation=True, - add_special_tokens=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - attention_mask = text_inputs.attention_mask - untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids - - if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): - removed_text = self.tokenizer_3.batch_decode( - untruncated_ids[:, min(max_sequence_length, self.tokenizer_3.model_max_length) - 1 : -1] - ) - logger.warning( - "The following part of your input was truncated because `max_sequence_length` is set to " - f" {min(max_sequence_length, self.tokenizer_3.model_max_length)} tokens: {removed_text}" - ) - - prompt_embeds = self.text_encoder_3( - text_input_ids.to(self.device), attention_mask=attention_mask.to(self.device) - )[0] - prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device) - return prompt_embeds - - def _get_clip_prompt_embeds( - self, - tokenizer, - text_encoder, - prompt: str | list[str], - max_sequence_length: int = 128, - dtype: torch.dtype | None = None, - ): - dtype = dtype or text_encoder.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - - model_max_length = self.tokenizer.model_max_length - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=min(max_sequence_length, model_max_length), - truncation=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids - if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): - removed_text = tokenizer.batch_decode(untruncated_ids[:, model_max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {model_max_length} tokens: {removed_text}" - ) - prompt_embeds = text_encoder(text_input_ids.to(self.device), output_hidden_states=True) - - # Use pooled output of CLIPTextModel - prompt_embeds = prompt_embeds[0] - prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device) - return prompt_embeds - - def _get_llama3_prompt_embeds( - self, - prompt: str | list[str] = None, - max_sequence_length: int = 128, - dtype: torch.dtype | None = None, - ): - dtype = dtype or self.text_encoder_4.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - - text_inputs = self.tokenizer_4( - prompt, - padding="max_length", - max_length=min(max_sequence_length, self.tokenizer_4.model_max_length), - truncation=True, - add_special_tokens=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - attention_mask = text_inputs.attention_mask - untruncated_ids = self.tokenizer_4(prompt, padding="longest", return_tensors="pt").input_ids - - if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): - removed_text = self.tokenizer_4.batch_decode( - untruncated_ids[:, min(max_sequence_length, self.tokenizer_4.model_max_length) - 1 : -1] - ) - logger.warning( - "The following part of your input was truncated because `max_sequence_length` is set to " - f" {min(max_sequence_length, self.tokenizer_4.model_max_length)} tokens: {removed_text}" - ) - - outputs = self.text_encoder_4( - text_input_ids.to(self.device), - attention_mask=attention_mask.to(self.device), - output_hidden_states=True, - ) - - prompt_embeds = outputs.hidden_states[1:] - prompt_embeds = torch.stack(prompt_embeds, dim=0) - return prompt_embeds - - def encode_prompt( - self, - prompt: str | list[str] | None = None, - prompt_2: str | list[str] | None = None, - prompt_3: str | list[str] | None = None, - prompt_4: str | list[str] | None = None, - dtype: torch.dtype | None = None, - num_images_per_prompt: int = 1, - do_classifier_free_guidance: bool = True, - negative_prompt: str | list[str] | None = None, - negative_prompt_2: str | list[str] | None = None, - negative_prompt_3: str | list[str] | None = None, - negative_prompt_4: str | list[str] | None = None, - prompt_embeds_t5: list[torch.FloatTensor] | None = None, - prompt_embeds_llama3: list[torch.FloatTensor] | None = None, - negative_prompt_embeds_t5: list[torch.FloatTensor] | None = None, - negative_prompt_embeds_llama3: list[torch.FloatTensor] | None = None, - pooled_prompt_embeds: torch.FloatTensor | None = None, - negative_pooled_prompt_embeds: torch.FloatTensor | None = None, - max_sequence_length: int = 128, - lora_scale: float | None = None, - ): - prompt = [prompt] if isinstance(prompt, str) else prompt - if prompt is not None: - batch_size = len(prompt) - else: - batch_size = pooled_prompt_embeds.shape[0] - - if pooled_prompt_embeds is None: - pooled_prompt_embeds_1 = self._get_clip_prompt_embeds( - self.tokenizer, self.text_encoder, prompt, max_sequence_length, dtype - ) - - if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt - - if len(negative_prompt) > 1 and len(negative_prompt) != batch_size: - raise ValueError(f"negative_prompt must be of length 1 or {batch_size}") - - negative_pooled_prompt_embeds_1 = self._get_clip_prompt_embeds( - self.tokenizer, self.text_encoder, negative_prompt, max_sequence_length, dtype - ) - - if negative_pooled_prompt_embeds_1.shape[0] == 1 and batch_size > 1: - negative_pooled_prompt_embeds_1 = negative_pooled_prompt_embeds_1.repeat(batch_size, 1) - - if pooled_prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - if len(prompt_2) > 1 and len(prompt_2) != batch_size: - raise ValueError(f"prompt_2 must be of length 1 or {batch_size}") - - pooled_prompt_embeds_2 = self._get_clip_prompt_embeds( - self.tokenizer_2, self.text_encoder_2, prompt_2, max_sequence_length, dtype - ) - - if pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1: - pooled_prompt_embeds_2 = pooled_prompt_embeds_2.repeat(batch_size, 1) - - if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: - negative_prompt_2 = negative_prompt_2 or negative_prompt - negative_prompt_2 = [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 - - if len(negative_prompt_2) > 1 and len(negative_prompt_2) != batch_size: - raise ValueError(f"negative_prompt_2 must be of length 1 or {batch_size}") - - negative_pooled_prompt_embeds_2 = self._get_clip_prompt_embeds( - self.tokenizer_2, self.text_encoder_2, negative_prompt_2, max_sequence_length, dtype - ) - - if negative_pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1: - negative_pooled_prompt_embeds_2 = negative_pooled_prompt_embeds_2.repeat(batch_size, 1) - - if pooled_prompt_embeds is None: - pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_1, pooled_prompt_embeds_2], dim=-1) - - if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: - negative_pooled_prompt_embeds = torch.cat( - [negative_pooled_prompt_embeds_1, negative_pooled_prompt_embeds_2], dim=-1 - ) - - if prompt_embeds_t5 is None: - prompt_3 = prompt_3 or prompt - prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3 - - if len(prompt_3) > 1 and len(prompt_3) != batch_size: - raise ValueError(f"prompt_3 must be of length 1 or {batch_size}") - - prompt_embeds_t5 = self._get_t5_prompt_embeds(prompt_3, max_sequence_length, dtype) - - if prompt_embeds_t5.shape[0] == 1 and batch_size > 1: - prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1) - - if do_classifier_free_guidance and negative_prompt_embeds_t5 is None: - negative_prompt_3 = negative_prompt_3 or negative_prompt - negative_prompt_3 = [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3 - - if len(negative_prompt_3) > 1 and len(negative_prompt_3) != batch_size: - raise ValueError(f"negative_prompt_3 must be of length 1 or {batch_size}") - - negative_prompt_embeds_t5 = self._get_t5_prompt_embeds(negative_prompt_3, max_sequence_length, dtype) - - if negative_prompt_embeds_t5.shape[0] == 1 and batch_size > 1: - negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1) - - if prompt_embeds_llama3 is None: - prompt_4 = prompt_4 or prompt - prompt_4 = [prompt_4] if isinstance(prompt_4, str) else prompt_4 - - if len(prompt_4) > 1 and len(prompt_4) != batch_size: - raise ValueError(f"prompt_4 must be of length 1 or {batch_size}") - - prompt_embeds_llama3 = self._get_llama3_prompt_embeds(prompt_4, max_sequence_length, dtype) - - if prompt_embeds_llama3.shape[0] == 1 and batch_size > 1: - prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1) - - if do_classifier_free_guidance and negative_prompt_embeds_llama3 is None: - negative_prompt_4 = negative_prompt_4 or negative_prompt - negative_prompt_4 = [negative_prompt_4] if isinstance(negative_prompt_4, str) else negative_prompt_4 - - if len(negative_prompt_4) > 1 and len(negative_prompt_4) != batch_size: - raise ValueError(f"negative_prompt_4 must be of length 1 or {batch_size}") - - negative_prompt_embeds_llama3 = self._get_llama3_prompt_embeds( - negative_prompt_4, max_sequence_length, dtype - ) - - if negative_prompt_embeds_llama3.shape[0] == 1 and batch_size > 1: - negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1) - - # duplicate pooled_prompt_embeds for each generation per prompt - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt) - pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - # duplicate t5_prompt_embeds for batch_size and num_images_per_prompt - bs_embed, seq_len, _ = prompt_embeds_t5.shape - if bs_embed == 1 and batch_size > 1: - prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1) - elif bs_embed > 1 and bs_embed != batch_size: - raise ValueError(f"cannot duplicate prompt_embeds_t5 of batch size {bs_embed}") - prompt_embeds_t5 = prompt_embeds_t5.repeat(1, num_images_per_prompt, 1) - prompt_embeds_t5 = prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1) - - # duplicate llama3_prompt_embeds for batch_size and num_images_per_prompt - _, bs_embed, seq_len, dim = prompt_embeds_llama3.shape - if bs_embed == 1 and batch_size > 1: - prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1) - elif bs_embed > 1 and bs_embed != batch_size: - raise ValueError(f"cannot duplicate prompt_embeds_llama3 of batch size {bs_embed}") - prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1) - prompt_embeds_llama3 = prompt_embeds_llama3.view(-1, batch_size * num_images_per_prompt, seq_len, dim) - - if do_classifier_free_guidance: - # duplicate negative_pooled_prompt_embeds for batch_size and num_images_per_prompt - bs_embed, seq_len = negative_pooled_prompt_embeds.shape - if bs_embed == 1 and batch_size > 1: - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(batch_size, 1) - elif bs_embed > 1 and bs_embed != batch_size: - raise ValueError(f"cannot duplicate negative_pooled_prompt_embeds of batch size {bs_embed}") - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt) - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - # duplicate negative_t5_prompt_embeds for batch_size and num_images_per_prompt - bs_embed, seq_len, _ = negative_prompt_embeds_t5.shape - if bs_embed == 1 and batch_size > 1: - negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1) - elif bs_embed > 1 and bs_embed != batch_size: - raise ValueError(f"cannot duplicate negative_prompt_embeds_t5 of batch size {bs_embed}") - negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds_t5 = negative_prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1) - - # duplicate negative_prompt_embeds_llama3 for batch_size and num_images_per_prompt - _, bs_embed, seq_len, dim = negative_prompt_embeds_llama3.shape - if bs_embed == 1 and batch_size > 1: - negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1) - elif bs_embed > 1 and bs_embed != batch_size: - raise ValueError(f"cannot duplicate negative_prompt_embeds_llama3 of batch size {bs_embed}") - negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1) - negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.view( - -1, batch_size * num_images_per_prompt, seq_len, dim - ) - - return ( - prompt_embeds_t5, - negative_prompt_embeds_t5, - prompt_embeds_llama3, - negative_prompt_embeds_llama3, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) - - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = ( - f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this" - " method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - ) - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = ( - f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this" - " method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - ) - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = ( - f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}`" - " is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - ) - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = ( - f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is" - " deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - ) - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - generator, - latents=None, - ): - # VAE applies 8x compression on images but we must also account for packing which requires - # latent height and width to be divisible by 2. - height = 2 * (int(height) // (self.vae_scale_factor * 2)) - width = 2 * (int(width) // (self.vae_scale_factor * 2)) - - shape = (batch_size, num_channels_latents, height, width) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=self.device, dtype=dtype) - else: - if latents.shape != shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") - latents = latents.to(self.device) - return latents - - def prepare_timesteps(self, num_inference_steps, sigmas, image_seq_len): - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas - mu = calculate_shift( - image_seq_len, - self.scheduler.config.get("base_image_seq_len", 256), - self.scheduler.config.get("max_image_seq_len", 4096), - self.scheduler.config.get("base_shift", 0.5), - self.scheduler.config.get("max_shift", 1.15), - ) - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, - num_inference_steps, - self.device, - sigmas=sigmas, - mu=mu, - ) - return timesteps, num_inference_steps - - def _extract_prompts(self, prompts): - """Extract prompt and negative_prompt from OmniPromptType list.""" - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None - if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts): - negative_prompt = None - elif prompts: - negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts] - else: - negative_prompt = None - return prompt, negative_prompt - - def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: bool): - if get_classifier_free_guidance_world_size() == 1: - return True - - if true_cfg_scale <= 1: - logger.warning("CFG parallel is NOT working correctly when true_cfg_scale <= 1.") - return False - - if not has_neg_prompt: - logger.warning( - "CFG parallel is NOT working correctly when there is no negative prompt or negative prompt embeddings." - ) - return False - return True - - def diffuse( - self, - prompt_embeds_t5: torch.Tensor, - prompt_embeds_llama3: torch.Tensor, - pooled_prompt_embeds: torch.Tensor, - latents: torch.Tensor, - timesteps: torch.Tensor, - do_true_cfg: bool, - ) -> torch.Tensor: - with self.progress_bar(total=len(timesteps)) as pbar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - noise_pred = self.transformer( - hidden_states=latent_model_input, - timesteps=timestep, - encoder_hidden_states_t5=prompt_embeds_t5, - encoder_hidden_states_llama3=prompt_embeds_llama3, - pooled_embeds=pooled_prompt_embeds, - return_dict=False, - )[0] - noise_pred = -noise_pred - - # TODO: Modify CFG guidance - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) - - pbar.update() - - return latents - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def do_classifier_free_guidance(self): - return self._guidance_scale > 1 - - @property - def attention_kwargs(self): - return self._attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @property - def interrupt(self): - return self._interrupt - - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | list[str] = None, - prompt_2: str | list[str] | None = None, - prompt_3: str | list[str] | None = None, - prompt_4: str | list[str] | None = None, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_scale: float = 5.0, - negative_prompt: str | list[str] | None = None, - negative_prompt_2: str | list[str] | None = None, - negative_prompt_3: str | list[str] | None = None, - negative_prompt_4: str | list[str] | None = None, - num_images_per_prompt: int | None = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.FloatTensor | None = None, - prompt_embeds_t5: torch.FloatTensor | None = None, - prompt_embeds_llama3: torch.FloatTensor | None = None, - negative_prompt_embeds_t5: torch.FloatTensor | None = None, - negative_prompt_embeds_llama3: torch.FloatTensor | None = None, - pooled_prompt_embeds: torch.FloatTensor | None = None, - negative_pooled_prompt_embeds: torch.FloatTensor | None = None, - output_type: str | None = "pil", - return_dict: bool = True, - attention_kwargs: dict[str, Any] | None = None, - callback_on_step_end: Callable[[int, int], None] | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 128, - **kwargs, - ): - extracted_prompt, negative_prompt = self._extract_prompts(req.prompts) - prompt = extracted_prompt or prompt - - height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor - width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor - - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - sigmas = req.sampling_params.sigmas or sigmas - max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - generator = req.sampling_params.generator or generator - true_cfg_scale = req.sampling_params.true_cfg_scale or guidance_scale - if req.sampling_params.guidance_scale_provided: - guidance_scale = req.sampling_params.guidance_scale - num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt - ) - - prompt_embeds = kwargs.get("prompt_embeds", None) - negative_prompt_embeds = kwargs.get("negative_prompt_embeds", None) - - if prompt_embeds is not None: - deprecation_message = ( - "The `prompt_embeds` argument is deprecated." - " Please use `prompt_embeds_t5` and `prompt_embeds_llama3` instead." - ) - deprecate("prompt_embeds", "0.35.0", deprecation_message) - prompt_embeds_t5 = prompt_embeds[0] - prompt_embeds_llama3 = prompt_embeds[1] - - if negative_prompt_embeds is not None: - deprecation_message = ( - "The `negative_prompt_embeds` argument is deprecated." - "Please use `negative_prompt_embeds_t5` and `negative_prompt_embeds_llama3` instead." - ) - deprecate("negative_prompt_embeds", "0.35.0", deprecation_message) - negative_prompt_embeds_t5 = negative_prompt_embeds[0] - negative_prompt_embeds_llama3 = negative_prompt_embeds[1] - - division = self.vae_scale_factor * 2 - S_max = (self.default_sample_size * self.vae_scale_factor) ** 2 - scale = S_max / (width * height) - scale = math.sqrt(scale) - width, height = int(width * scale // division * division), int(height * scale // division * division) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - prompt_3, - prompt_4, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - negative_prompt_4=negative_prompt_4, - prompt_embeds_t5=prompt_embeds_t5, - prompt_embeds_llama3=prompt_embeds_llama3, - negative_prompt_embeds_t5=negative_prompt_embeds_t5, - negative_prompt_embeds_llama3=negative_prompt_embeds_llama3, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._interrupt = False - - # 2. Define call parameters - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - elif pooled_prompt_embeds is not None: - batch_size = pooled_prompt_embeds.shape[0] - - # TODO: CFG guidance configuration - has_neg_prompt = negative_prompt is not None or ( - negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None - ) - do_true_cfg = true_cfg_scale > 1 and has_neg_prompt - self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt) - - # 3. Encode prompt - lora_scale = self.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None - ( - prompt_embeds_t5, - negative_prompt_embeds_t5, - prompt_embeds_llama3, - negative_prompt_embeds_llama3, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_3=prompt_3, - prompt_4=prompt_4, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - negative_prompt_4=negative_prompt_4, - do_classifier_free_guidance=self.do_classifier_free_guidance, - prompt_embeds_t5=prompt_embeds_t5, - prompt_embeds_llama3=prompt_embeds_llama3, - negative_prompt_embeds_t5=negative_prompt_embeds_t5, - negative_prompt_embeds_llama3=negative_prompt_embeds_llama3, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - - if self.do_classifier_free_guidance: - prompt_embeds_t5 = torch.cat([negative_prompt_embeds_t5, prompt_embeds_t5], dim=0) - prompt_embeds_llama3 = torch.cat([negative_prompt_embeds_llama3, prompt_embeds_llama3], dim=1) - pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) - - # 4. Prepare latent variables - num_channels_latents = self.transformer.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - pooled_prompt_embeds.dtype, - generator, - latents, - ) - - # 5. Prepare timesteps - timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1]) - self._num_timesteps = len(timesteps) - - # 6. Denoising loop - latents = self.diffuse( - prompt_embeds_t5, - prompt_embeds_llama3, - pooled_prompt_embeds, - latents, - timesteps, - do_true_cfg, - ) - - if output_type == "latent": - image = latents - - else: - latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor - - image = self.vae.decode(latents, return_dict=False)[0] - - return DiffusionOutput( - output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None - ) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) From 63ba11b8b1f840b8b5fd7ee8232d7d9cd0a424ff Mon Sep 17 00:00:00 2001 From: yangyonggit Date: Fri, 19 Jun 2026 02:38:32 +0200 Subject: [PATCH 12/40] [Perf] Stacked audio ops: batched matmul and embedding gather in MOSS-TTS talker (#4230) Signed-off-by: leo.yang --- tests/e2e/offline_inference/test_moss_tts.py | 327 +++++++++--------- .../test_moss_tts_realtime.py | 290 ++++++++++++++++ .../moss_tts/modeling_moss_tts_talker.py | 83 ++++- 3 files changed, 534 insertions(+), 166 deletions(-) create mode 100644 tests/e2e/offline_inference/test_moss_tts_realtime.py diff --git a/tests/e2e/offline_inference/test_moss_tts.py b/tests/e2e/offline_inference/test_moss_tts.py index 505db1af7..19a8ac6d6 100644 --- a/tests/e2e/offline_inference/test_moss_tts.py +++ b/tests/e2e/offline_inference/test_moss_tts.py @@ -1,77 +1,97 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""E2E offline inference tests for MOSS-TTS two-stage pipeline. +"""E2E offline inference tests for MOSS-VoiceGenerator (MossTTSDelayModel, 1.7B). -Tests the MossTTSDelayModel path using MOSS-VoiceGenerator (1.7B) — the -smallest MossTTSDelayModel variant — so the test runs on a single A10G or L4 -without requiring an 80 GB GPU. +Uses the standard omni_runner + pytestmark pattern (one module-scoped engine +per file, skill invariant I4). The realtime variant lives in the sibling file +test_moss_tts_realtime.py. -A second test class covers MOSS-TTS-Realtime (1.7B, MossTTSRealtime) to -exercise the local-transformer generation path. +MOSS-VoiceGenerator synthesizes speech from a text *instruction* that describes +the desired voice (e.g. "a warm female voice with an American accent"). It does +NOT accept a reference audio clip. The request format requires running +AutoProcessor.from_pretrained once per call to encode the (text, instruction) +pair into the (prompt_token_ids, codes.ref) grid that the MossTTSDelayModel +talker expects — same as examples/offline_inference/text_to_speech/moss_tts/ +end2end.py:_build_unified_codes. -Both classes use a module-scoped engine so the model is loaded once per test -module execution. Do NOT add a second engine fixture in this file — that -triggers mid-module teardown races (see skill invariant I4). +No determinism test: VoiceGenerator produces variable-length output even with +a fixed seed; waveform length reproducibility is not guaranteed. """ from __future__ import annotations +import gc import os -import urllib.request import pytest import torch from vllm import SamplingParams from tests.helpers.mark import hardware_test -from vllm_omni import Omni +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config # --------------------------------------------------------------------------- -# Shared constants +# Constants # --------------------------------------------------------------------------- -# Selected by the nightly "TTS · Function Test" (-m "full_model and L4 and tts"). -# Without the tts + run-level marks these tests are collected but never run. -# TODO: Fix this test -# pytestmark = [pytest.mark.full_model, pytest.mark.tts] +MODEL = "OpenMOSS-Team/MOSS-VoiceGenerator" +SAMPLE_RATE = 24_000 -SAMPLE_RATE = 24_000 # All MOSS-TTS full variants output 24 kHz - -REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav" - -_DEFAULT_SAMPLING = SamplingParams( +# Stage 0 = AR talker; max_tokens=512 keeps tests fast (~20 s of audio). +# Stage 1 = codec decoder; greedy, large context to hold all codec tokens. +_STAGE0_PARAMS = SamplingParams( temperature=1.7, top_p=0.8, top_k=25, - max_tokens=512, # short for tests (~20 s at 24 kHz, ~12.5 tokens/s) + max_tokens=512, seed=42, detokenize=False, ) +_STAGE1_PARAMS = SamplingParams( + temperature=0.0, + top_p=1.0, + top_k=-1, + max_tokens=65536, + seed=42, + detokenize=False, +) +_SAMPLING = [_STAGE0_PARAMS, _STAGE1_PARAMS] # --------------------------------------------------------------------------- -# Session fixtures +# Deploy config # --------------------------------------------------------------------------- -@pytest.fixture(scope="session") -def ref_audio_path(tmp_path_factory) -> str: - """Download the upstream reference clip once per session. +def _get_test_config() -> str: + """Derive a CI-friendly config from moss_voice_generator.yaml. - Set ``MOSS_TTS_SKIP_ON_NET_FAIL=1`` to skip in air-gapped environments. + Reduces Stage 0 gpu_memory_utilization from 0.60 → 0.45 to leave headroom + for Stage 1 (0.12) on a shared L4/A10G. max_num_seqs=1 keeps per-test + peak memory predictable. """ - cache_dir = tmp_path_factory.mktemp("moss_tts_ref") - target = cache_dir / "zh_1.wav" - try: - with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp: - target.write_bytes(resp.read()) - except Exception as exc: - msg = f"Cannot fetch reference clip {REF_AUDIO_URL}: {exc}" - if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): - pytest.skip(msg) - pytest.fail(msg) - if not target.exists() or target.stat().st_size == 0: - pytest.fail(f"Reference clip empty after download: {target}") - return str(target) + return modify_stage_config( + get_deploy_config_path("moss_voice_generator.yaml"), + updates={ + "stages": { + 0: { + "max_num_seqs": 1, + "gpu_memory_utilization": 0.45, + }, + }, + }, + ) + + +# --------------------------------------------------------------------------- +# pytestmark — one engine for the whole module +# --------------------------------------------------------------------------- + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.tts, + pytest.mark.parametrize("omni_runner", [(MODEL, _get_test_config())], indirect=True), +] # --------------------------------------------------------------------------- @@ -79,147 +99,142 @@ def ref_audio_path(tmp_path_factory) -> str: # --------------------------------------------------------------------------- -def _build_request(text: str, ref_audio_path: str, seed: int = 42) -> dict: +def _build_request(text: str, instruction: str) -> dict: + """Build a VoiceGenerator request using AutoProcessor. + + Calls proc.build_user_message + proc() to produce the unified-codes grid + (L, 1+n_vq) and splits it into prompt_token_ids + codes.ref, exactly as + end2end.py:_build_unified_codes does for the voice_generator variant. + The processor (including its CPU-resident audio_tokenizer) is freed before + returning so it does not compete with the running engine. + + Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments where + the HF snapshot is not cached. + """ + from transformers import AutoProcessor + + try: + proc = AutoProcessor.from_pretrained(MODEL, trust_remote_code=True) + except Exception as exc: + msg = f"Cannot load AutoProcessor for {MODEL}: {exc}" + if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): + pytest.skip(msg) + pytest.fail(msg) + + user_msg = proc.build_user_message(text=text, instruction=instruction) + batch = proc(conversations=[[user_msg]], mode="generation") + unified = batch["input_ids"][0] # (L, 1+n_vq) + text_ids = unified[:, 0].tolist() + audio_codes = unified[:, 1:].contiguous().to(torch.int64) # (L, n_vq) + del proc + gc.collect() + return { - "prompt": "<|im_start|>assistant\n", - "additional_information": { - "task_type": ["voice_clone"], - "text": [text], - "mode": ["voice_clone"], - "prompt_audio_path": [ref_audio_path], - "seed": [seed], - }, + "prompt_token_ids": text_ids, + "additional_information": {"codes": {"ref": audio_codes}}, } -def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]: - """Run one request; return (waveform_cpu, sample_rate).""" - for stage_outputs in omni.generate(request, _DEFAULT_SAMPLING): - for req_output in stage_outputs.request_output: - mm = req_output.outputs[0].multimodal_output - assert mm is not None, "Expected non-None multimodal_output" - # Never use ``a or b`` on tensor values — a multi-element tensor - # raises on truthiness. Check ``is None`` explicitly and handle the - # pre-consolidation list form. - audio = mm.get("audio") - if audio is None: - audio = mm.get("model_outputs") - sr = mm.get("sr") - assert audio is not None, "Expected 'audio' or 'model_outputs' in multimodal_output" - if isinstance(audio, list): - audio = torch.cat( - [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0], - dim=0, - ) - assert isinstance(audio, torch.Tensor), f"Expected Tensor, got {type(audio)}" - return audio.reshape(-1).cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE - raise AssertionError("No stage outputs received") +def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor, int]: + """Run one request and return (waveform_cpu, sample_rate). + + Iterates all OmniRequestOutputs from generate() via the top-level + multimodal_output property (which aggregates from completion outputs). + AR-stage outputs have empty multimodal_output and are silently skipped. + With async_chunk=True and FINAL_ONLY, the codec stage consolidates its + chunks into a single tensor before yielding — but we handle list fallback + for robustness. + """ + chunks: list[torch.Tensor] = [] + sr_final = SAMPLE_RATE + for out in omni_runner.omni.generate(request, _SAMPLING): + mm = out.multimodal_output + if not mm: + continue + audio = mm.get("audio") + if audio is None: + audio = mm.get("model_outputs") + if audio is None: + continue + sr = mm.get("sr") + if sr is not None: + sr_final = int(sr.item() if hasattr(sr, "item") else sr) + if isinstance(audio, list): + audio = torch.cat( + [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0], + dim=0, + ) + if isinstance(audio, torch.Tensor) and audio.numel() > 0: + chunks.append(audio.reshape(-1).cpu()) + if not chunks: + raise AssertionError("No audio output received from generate()") + return torch.cat(chunks, dim=0), sr_final # --------------------------------------------------------------------------- -# MossTTSDelayModel — MOSS-VoiceGenerator 1.7B -# (tests the delay-pattern generation path on a small GPU) +# Tests # --------------------------------------------------------------------------- -_DELAY_MODEL = "OpenMOSS-Team/MOSS-VoiceGenerator" - -@pytest.fixture(scope="module") -def delay_engine(): - return Omni(model=_DELAY_MODEL, stage_init_timeout=300) - - -@hardware_test(res={"cuda": "L4"}) -def test_moss_tts_delay_english(delay_engine, ref_audio_path): - """MossTTSDelayModel: English voice_clone produces non-empty 24 kHz audio.""" - req = _build_request("Hello, this is a MOSS-TTS voice cloning test.", ref_audio_path) - audio, sr = _collect_audio(delay_engine, req) +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_moss_tts_english(omni_runner: OmniRunner) -> None: + """VoiceGenerator: English instruction produces non-empty 24 kHz audio.""" + req = _build_request( + "Hello, this is a MOSS voice design test.", + "a warm female voice with an American accent", + ) + audio, sr = _collect_audio(omni_runner, req) assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}" assert audio.numel() > 0, "Audio tensor is empty" assert not torch.all(audio == 0), "Audio is silence" -@hardware_test(res={"cuda": "L4"}) -def test_moss_tts_delay_chinese(delay_engine, ref_audio_path): - """MossTTSDelayModel: Chinese input produces non-empty audio.""" - req = _build_request("你好，这是语音合成测试。", ref_audio_path) - audio, sr = _collect_audio(delay_engine, req) +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_moss_tts_chinese(omni_runner: OmniRunner) -> None: + """VoiceGenerator: Chinese input produces non-empty audio.""" + req = _build_request( + "你好，这是语音合成测试。", + "一个清晰温暖的女声", + ) + audio, sr = _collect_audio(omni_runner, req) assert sr == SAMPLE_RATE assert audio.numel() > 0 assert not torch.all(audio == 0) -@hardware_test(res={"cuda": "L4"}) -def test_moss_tts_delay_deterministic(delay_engine, ref_audio_path): - """MossTTSDelayModel: same seed yields identical waveforms.""" - req = _build_request("Reproducibility check.", ref_audio_path, seed=99) - audio1, _ = _collect_audio(delay_engine, req) - audio2, _ = _collect_audio(delay_engine, req) - - assert audio1.shape == audio2.shape, "Shapes differ across identical seeds" - assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms differ across identical seeds" - - -@hardware_test(res={"cuda": "L4"}) -def test_moss_tts_delay_batch(delay_engine, ref_audio_path): - """MossTTSDelayModel: batch of two requests each returns non-empty audio.""" +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_moss_tts_batch(omni_runner: OmniRunner) -> None: + """VoiceGenerator: batch of two requests each returns non-empty audio.""" requests = [ - _build_request("First sentence.", ref_audio_path), - _build_request("Second sentence.", ref_audio_path), + _build_request("First sentence.", "a warm female voice"), + _build_request("Second sentence.", "a young male voice"), ] results: list[torch.Tensor] = [] - for stage_outputs in delay_engine.generate(requests, [_DEFAULT_SAMPLING] * 2): - for req_output in stage_outputs.request_output: - mm = req_output.outputs[0].multimodal_output - assert mm is not None - audio = mm.get("audio") - if audio is None: - audio = mm.get("model_outputs") - assert audio is not None - if isinstance(audio, list): - audio = torch.cat( - [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0], - dim=0, - ) + for out in omni_runner.omni.generate(requests, _SAMPLING): + mm = out.multimodal_output + if not mm: + continue + audio = mm.get("audio") + if audio is None: + audio = mm.get("model_outputs") + if audio is None: + continue + if isinstance(audio, list): + audio = torch.cat( + [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0], + dim=0, + ) + if isinstance(audio, torch.Tensor) and audio.numel() > 0: results.append(audio.reshape(-1).cpu()) - assert len(results) == 2, f"Expected 2 outputs, got {len(results)}" + # With async_chunk streaming + FINAL_ONLY, each request produces one + # consolidated output. The >= 2 guard handles backward-compatibility if + # consolidation behavior changes. + assert len(results) >= 2, f"Expected at least 2 audio outputs, got {len(results)}" for i, audio in enumerate(results): - assert audio.numel() > 0, f"Audio[{i}] is empty" - - -# --------------------------------------------------------------------------- -# MossTTSRealtime — MOSS-TTS-Realtime 1.7B -# (tests the local-transformer generation path) -# --------------------------------------------------------------------------- - -_REALTIME_MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime" - - -@pytest.fixture(scope="module") -def realtime_engine(): - return Omni(model=_REALTIME_MODEL, stage_init_timeout=300) - - -@hardware_test(res={"cuda": "L4"}) -def test_moss_tts_realtime_english(realtime_engine, ref_audio_path): - """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio.""" - req = _build_request("This is a real-time TTS streaming test.", ref_audio_path) - audio, sr = _collect_audio(realtime_engine, req) - - assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}" - assert audio.numel() > 0, "Audio tensor is empty" - assert not torch.all(audio == 0), "Audio is silence" - - -@hardware_test(res={"cuda": "L4"}) -def test_moss_tts_realtime_deterministic(realtime_engine, ref_audio_path): - """MossTTSRealtime: same seed yields identical waveforms.""" - req = _build_request("Determinism check for realtime model.", ref_audio_path, seed=7) - audio1, _ = _collect_audio(realtime_engine, req) - audio2, _ = _collect_audio(realtime_engine, req) - - assert audio1.shape == audio2.shape - assert torch.allclose(audio1, audio2, atol=1e-4) + assert audio.numel() > 0, f"Audio chunk[{i}] is empty" diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime.py b/tests/e2e/offline_inference/test_moss_tts_realtime.py new file mode 100644 index 000000000..540c8750e --- /dev/null +++ b/tests/e2e/offline_inference/test_moss_tts_realtime.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E offline inference tests for MOSS-TTS-Realtime (MossTTSRealtime, 1.7B). + +Uses the standard omni_runner + pytestmark pattern (one module-scoped engine +per file, skill invariant I4). The delay-model variant lives in the sibling +file test_moss_tts.py. + +The request format mirrors end2end.py:_build_realtime_prompt exactly: + 1. snapshot_download to locate the MossTTSRealtimeProcessor module. + 2. MOSS-Audio-Tokenizer (separate HF repo) to encode the reference clip. + 3. Build the (L, 17) int grid; split into prompt_token_ids + codes.ref + + ids.all (remaining text tokens fed one-per-step during decode). + +Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments. + +No determinism test: MossTTSRealtime uses a local depth transformer that is +sensitive to async scheduling timing; waveform reproducibility is not +guaranteed across back-to-back calls. +""" + +from __future__ import annotations + +import gc +import importlib.util +import os +import sys +import urllib.request +from pathlib import Path + +import numpy as np +import pytest +import soundfile as sf +import torch +from vllm import SamplingParams + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime" +SAMPLE_RATE = 24_000 + +REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav" + +_STAGE0_PARAMS = SamplingParams( + temperature=1.0, + top_p=0.95, + top_k=50, + max_tokens=512, + seed=42, + detokenize=False, +) +_STAGE1_PARAMS = SamplingParams( + temperature=0.0, + top_p=1.0, + top_k=-1, + max_tokens=65536, + seed=42, + detokenize=False, +) +_SAMPLING = [_STAGE0_PARAMS, _STAGE1_PARAMS] + +# --------------------------------------------------------------------------- +# Deploy config +# --------------------------------------------------------------------------- + + +def _get_test_config() -> str: + """Derive a CI-friendly config from moss_tts_realtime.yaml. + + Reduces Stage 0 gpu_memory_utilization from 0.60 → 0.45 to leave headroom + for Stage 1 (0.12) on a shared L4/A10G. max_num_seqs=1 keeps per-test + peak memory predictable. + """ + return modify_stage_config( + get_deploy_config_path("moss_tts_realtime.yaml"), + updates={ + "stages": { + 0: { + "max_num_seqs": 1, + "gpu_memory_utilization": 0.45, + }, + }, + }, + ) + + +# --------------------------------------------------------------------------- +# pytestmark — one engine for the whole module +# --------------------------------------------------------------------------- + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.tts, + pytest.mark.parametrize("omni_runner", [(MODEL, _get_test_config())], indirect=True), +] + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def ref_audio_path(tmp_path_factory) -> str: + """Download the upstream reference clip once per test module. + + Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments. + """ + target = tmp_path_factory.mktemp("moss_tts_realtime_ref") / "zh_1.wav" + try: + with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp: + target.write_bytes(resp.read()) + except Exception as exc: + msg = f"Cannot fetch reference clip {REF_AUDIO_URL}: {exc}" + if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): + pytest.skip(msg) + pytest.fail(msg) + if not target.exists() or target.stat().st_size == 0: + pytest.fail(f"Reference clip empty after download: {target}") + return str(target) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _load_ref_audio(path: str, target_sr: int = 24000) -> torch.Tensor: + """Load and resample reference audio to (1, T) float32 at target_sr.""" + data, sr = sf.read(path, always_2d=True) # (T, C) + wav = torch.from_numpy(data.T.astype("float32")) # (C, T) + if wav.shape[0] > 1: + wav = wav.mean(dim=0, keepdim=True) + if sr != target_sr: + import torchaudio + + wav = torchaudio.functional.resample(wav, sr, target_sr) + return wav # (1, T) + + +def _build_request(ref_audio_path: str, text: str) -> dict: + """Build a MOSS-TTS-Realtime request. + + Mirrors end2end.py:_build_realtime_prompt exactly: + 1. Locate MossTTSRealtimeProcessor in the HF snapshot. + 2. Encode ref audio via MOSS-Audio-Tokenizer (separate repo). + 3. Build the (L, 17) int grid and return prompt_token_ids + + additional_information: {codes.ref, ids.all}. + + Frees the codec model before returning to avoid competing with the + running vllm engine for GPU/CPU memory. + """ + from huggingface_hub import snapshot_download + from transformers import AutoModel, AutoTokenizer + + # Step 1: locate the realtime processor module in the snapshot. + try: + snap_dir = Path(snapshot_download(repo_id=MODEL)) + except Exception as exc: + msg = f"Cannot locate snapshot for {MODEL}: {exc}" + if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): + pytest.skip(msg) + pytest.fail(msg) + + proc_module_path = snap_dir / "processing_mossttsrealtime.py" + if not proc_module_path.exists(): + pytest.fail(f"Realtime processor module missing at {proc_module_path}") + + mod_name = "_moss_tts_realtime_proc" + if mod_name not in sys.modules: + spec = importlib.util.spec_from_file_location(mod_name, proc_module_path) + mod = importlib.util.module_from_spec(spec) # type: ignore[arg-type] + sys.modules[mod_name] = mod + spec.loader.exec_module(mod) # type: ignore[union-attr] + else: + mod = sys.modules[mod_name] + + tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True) + realtime_processor = mod.MossTTSRealtimeProcessor(tokenizer=tokenizer) + + # Step 2: encode the reference audio clip via MOSS-Audio-Tokenizer. + try: + codec = AutoModel.from_pretrained("OpenMOSS-Team/MOSS-Audio-Tokenizer", trust_remote_code=True).to("cpu").eval() + except Exception as exc: + msg = f"Cannot load MOSS-Audio-Tokenizer: {exc}" + if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): + pytest.skip(msg) + pytest.fail(msg) + + ref_wav = _load_ref_audio(ref_audio_path) + wav_2d = ref_wav if ref_wav.dim() == 2 else ref_wav.unsqueeze(0) + with torch.no_grad(): + enc = codec.batch_encode([wav_2d.squeeze(0)], num_quantizers=16) + codes = enc.audio_codes[:, 0, : int(enc.audio_codes_lengths[0].item())] + audio_tokens = codes.transpose(0, 1).contiguous().cpu().numpy() # (T, 16) + del codec + gc.collect() + + # Step 3: build the (L, 17) prefill grid. + system_grid = realtime_processor.make_ensemble(prompt_audio_tokens=audio_tokens) + + assistant_header_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False) + assistant_grid = np.full((len(assistant_header_ids), 17), 1024, dtype=np.int64) + assistant_grid[:, 0] = assistant_header_ids + + text_ids_only = tokenizer.encode(text, add_special_tokens=False) + if not text_ids_only: + pytest.fail(f"Empty text after tokenization: {text!r}") + + PREFILL_MAX_TEXT = 12 + cur_len = min(len(text_ids_only), PREFILL_MAX_TEXT) + text_grid = np.full((cur_len, 17), 1024, dtype=np.int64) + text_grid[:, 0] = text_ids_only[:cur_len] + text_grid[-1, 1] = 1025 # audio_bos at the last prefilled text token + + grid = np.concatenate([system_grid, assistant_grid, text_grid], axis=0) + prompt_token_ids = grid[:, 0].tolist() + audio_codes_tensor = torch.from_numpy(grid[:, 1:].astype(np.int64).copy()) # (L, 16) + remaining_text_ids = list(text_ids_only[cur_len:]) + + additional_info: dict = {"codes": {"ref": audio_codes_tensor}} + if remaining_text_ids: + additional_info["ids"] = {"all": remaining_text_ids} + + return { + "prompt_token_ids": prompt_token_ids, + "additional_information": additional_info, + } + + +def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor, int]: + """Run one request and return (waveform_cpu, sample_rate).""" + chunks: list[torch.Tensor] = [] + sr_final = SAMPLE_RATE + for out in omni_runner.omni.generate(request, _SAMPLING): + mm = out.multimodal_output + if not mm: + continue + audio = mm.get("audio") + if audio is None: + audio = mm.get("model_outputs") + if audio is None: + continue + sr = mm.get("sr") + if sr is not None: + sr_final = int(sr.item() if hasattr(sr, "item") else sr) + if isinstance(audio, list): + audio = torch.cat( + [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0], + dim=0, + ) + if isinstance(audio, torch.Tensor) and audio.numel() > 0: + chunks.append(audio.reshape(-1).cpu()) + if not chunks: + raise AssertionError("No audio output received from generate()") + return torch.cat(chunks, dim=0), sr_final + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) -> None: + """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio.""" + req = _build_request(ref_audio_path, "This is a real-time TTS streaming test.") + audio, sr = _collect_audio(omni_runner, req) + + assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}" + assert audio.numel() > 0, "Audio tensor is empty" + assert not torch.all(audio == 0), "Audio is silence" + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_moss_tts_realtime_chinese(omni_runner: OmniRunner, ref_audio_path: str) -> None: + """MossTTSRealtime: Chinese input produces non-empty audio.""" + req = _build_request(ref_audio_path, "你好，这是语音合成测试。") + audio, sr = _collect_audio(omni_runner, req) + + assert sr == SAMPLE_RATE + assert audio.numel() > 0 + assert not torch.all(audio == 0) diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py index a03b4c1d1..d40e22014 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py @@ -153,6 +153,12 @@ class MossTTSDelayTalkerForGeneration(nn.Module): # `make_omni_output`, which runs immediately before the sampler). self._batch_state: list[dict[str, Any]] | None = None + # Stacked weight caches built after load_weights(). Avoids a Python + # loop over n_vq heads/embeddings at every decode step. + # shapes: (n_vq, audio_vocab+1, hidden_size) + self._stacked_audio_head_w: torch.Tensor | None = None + self._stacked_audio_emb_w: torch.Tensor | None = None + # ------------------------------------------------------------------ # vLLM-Omni hooks # ------------------------------------------------------------------ @@ -394,9 +400,15 @@ class MossTTSDelayTalkerForGeneration(nn.Module): chunk = ref_codes[ref_offset:end_off] if chunk.numel() > 0 and chunk.shape[0] == span_len: codes = chunk.to(device=device, dtype=torch.long).clamp_(0, self.audio_vocab_size) - audio_embed = torch.zeros_like(embeds) - for i, emb_layer in enumerate(self.audio_embeddings): - audio_embed = audio_embed + emb_layer(codes[:, i]) + if self._stacked_audio_emb_w is not None: + codes_nq = codes.t() # (n_vq, L) + audio_embed = self._stacked_audio_emb_w[ + torch.arange(self.n_vq, device=device)[:, None], codes_nq + ].sum(0) + else: + audio_embed = torch.zeros_like(embeds) + for i, emb_layer in enumerate(self.audio_embeddings): + audio_embed = audio_embed + emb_layer(codes[:, i]) embeds = embeds + audio_embed last_ref_row = codes[-1].detach() @@ -436,10 +448,17 @@ class MossTTSDelayTalkerForGeneration(nn.Module): audio_codes_buf = (info_dict.get("audio_codes", {}) or {}).get("current") if isinstance(audio_codes_buf, torch.Tensor) and audio_codes_buf.numel() == self.n_vq: codes = audio_codes_buf.to(device=device, dtype=torch.long) - audio_embed = torch.zeros_like(text_embed) - for i, emb_layer in enumerate(self.audio_embeddings): - code_i = codes[i].clamp(0, self.audio_vocab_size).reshape(1) - audio_embed = audio_embed + emb_layer(code_i) + if self._stacked_audio_emb_w is not None: + # Single gather: stacked_emb[i, codes[i]] for each i → (n_vq, H) → sum → (1, H) + n_vq_idx = torch.arange(self.n_vq, device=device) + audio_embed = self._stacked_audio_emb_w[n_vq_idx, codes.clamp(0, self.audio_vocab_size)].sum( + 0, keepdim=True + ) + else: + audio_embed = torch.zeros_like(text_embed) + for i, emb_layer in enumerate(self.audio_embeddings): + code_i = codes[i].clamp(0, self.audio_vocab_size).reshape(1) + audio_embed = audio_embed + emb_layer(code_i) combined = text_embed + audio_embed else: combined = text_embed @@ -501,17 +520,38 @@ class MossTTSDelayTalkerForGeneration(nn.Module): if not bool(sampling_mask.any()): return codes - # Use deploy YAML defaults. The audio sampler isn't routed through - # vLLM's sampler so we apply the upstream's audio_top_k locally. audio_top_k = 25 audio_temp = 1.7 + if self._stacked_audio_head_w is not None: + # Single batched matmul: (n_vq, V+1, H) @ (H,) → (n_vq, V+1). + # Replaces n_vq serial nn.Linear calls. + all_logits = self._stacked_audio_head_w @ last_h.reshape(-1) # (n_vq, V+1) + all_logits[:, -1] = float("-inf") # invalid sentinel + all_logits[:, self.audio_pad_code] = float("-inf") + + active_idx = sampling_mask.nonzero(as_tuple=True)[0] # (n_active,) + active_logits = all_logits[active_idx] # (n_active, V+1) + active_logits = active_logits / max(audio_temp, 1e-6) + if audio_top_k > 0 and audio_top_k < active_logits.shape[-1]: + top_vals, _ = torch.topk(active_logits, audio_top_k, dim=-1) + kth = top_vals[:, -1:].expand_as(active_logits) + active_logits = torch.where( + active_logits < kth, + torch.full_like(active_logits, float("-inf")), + active_logits, + ) + probs = torch.softmax(active_logits, dim=-1) + codes[active_idx] = torch.multinomial(probs, num_samples=1).squeeze(-1).long() + return codes + + # Fallback: per-head loop (used before load_weights() completes). for i in range(n_vq): if not bool(sampling_mask[i]): continue head = self.audio_heads[i] logits = head(last_h).reshape(-1) # (V,) - logits[..., -1] = float("-inf") # invalid sentinel + logits[..., -1] = float("-inf") logits[..., self.audio_pad_code] = float("-inf") sampled = self._sample_with_top_k(logits, audio_top_k, audio_temp) codes[i] = sampled.long() @@ -678,6 +718,14 @@ class MossTTSDelayTalkerForGeneration(nn.Module): for n in backbone_loaded: loaded.add(f"model.{n}") + # Build stacked weight caches for vectorised decode ops. + self._stacked_audio_head_w = torch.stack( + [h.weight.detach() for h in self.audio_heads], dim=0 + ) # (n_vq, audio_vocab+1, hidden_size) + self._stacked_audio_emb_w = torch.stack( + [e.weight.detach() for e in self.audio_embeddings], dim=0 + ) # (n_vq, audio_vocab+1, hidden_size) + return loaded @@ -765,6 +813,10 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): self.logits_processor = LogitsProcessor(self.text_vocab_size) self._batch_state: list[dict[str, Any]] | None = None + # Stacked weight cache built after load_weights(). + # shape: (n_vq, audio_vocab_size, hidden_size) + self._stacked_audio_emb_w: torch.Tensor | None = None + self.gpu_resident_buffer_keys: set[tuple[str, str]] = { ("audio_codes", "current"), ("audio_codes", "accumulated"), @@ -854,6 +906,11 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): embeds = self.embed_tokens[0](text_ids) if audio_codes is None: return embeds + if self._stacked_audio_emb_w is not None: + # Single gather: (n_vq, vocab, H)[arange[:,None], codes.T] → (n_vq, T, H) → sum → (T, H) + n_vq_idx = torch.arange(self.n_vq, device=audio_codes.device) + codes_t = audio_codes.t().clamp(0, self.audio_vocab_size - 1) # (n_vq, T) + return embeds + self._stacked_audio_emb_w[n_vq_idx.unsqueeze(1), codes_t].sum(0) for i in range(self.n_vq): col = audio_codes[:, i].clamp_(0, self.audio_vocab_size - 1) embeds = embeds + self.embed_tokens[i + 1](col) @@ -1115,6 +1172,12 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): for n in backbone_loaded: loaded.add(f"model.{n}") + # Build stacked embedding cache for vectorised decode ops. + # embed_tokens[0] is the text embedding; [1..n_vq] are audio codebooks. + self._stacked_audio_emb_w = torch.stack( + [self.embed_tokens[i + 1].weight.detach() for i in range(self.n_vq)], dim=0 + ) # (n_vq, audio_vocab_size, hidden_size) + logger.info( "[MossTTSRealtime] loaded %d/%d params; skipped=%d (first 5: %s)", len(loaded), From f05e03aea84d1e9a22a45ed0e995ffc2a1eea730 Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Fri, 19 Jun 2026 14:24:37 +0800 Subject: [PATCH 13/40] [Perf][DreamZero]dreamzero supports CUDA graph, torch.compile, and DiT Caching. (#4213) Signed-off-by: amy-why-3459 --- .../dreamzero/client_schedule.py | 129 +++++++++ .../dreamzero/export_prediction_video.py | 136 +++++---- tests/diffusion/cache/test_stepcache.py | 116 ++++++++ .../test_dreamzero_expansion.py | 2 - .../openai_api/test_openpi_connection.py | 41 +++ vllm_omni/deploy/dreamzero.yaml | 7 +- vllm_omni/deploy/dreamzero_tp1_cfg2.yaml | 7 +- vllm_omni/diffusion/cache/base.py | 2 + vllm_omni/diffusion/cache/selector.py | 13 +- .../diffusion/cache/stepcache/__init__.py | 66 +++++ .../diffusion/cache/stepcache/backend.py | 129 +++++++++ vllm_omni/diffusion/cache/stepcache/config.py | 51 ++++ vllm_omni/diffusion/cache/stepcache/state.py | 51 ++++ vllm_omni/diffusion/data.py | 9 + .../models/dreamzero/pipeline_dreamzero.py | 272 ++++++++++++++---- .../models/dreamzero/state_dreamzero.py | 99 +++++-- .../models/dreamzero/video_export_worker.py | 25 +- .../dreamzero/wan_vae_feat_cache_patch.py | 92 ++++++ .../worker/diffusion_model_runner.py | 24 +- vllm_omni/engine/async_omni_engine.py | 10 +- vllm_omni/entrypoints/cli/serve.py | 2 +- vllm_omni/entrypoints/openpi/connection.py | 98 +++++-- 22 files changed, 1204 insertions(+), 177 deletions(-) create mode 100644 examples/offline_inference/dreamzero/client_schedule.py create mode 100644 tests/diffusion/cache/test_stepcache.py create mode 100644 vllm_omni/diffusion/cache/stepcache/__init__.py create mode 100644 vllm_omni/diffusion/cache/stepcache/backend.py create mode 100644 vllm_omni/diffusion/cache/stepcache/config.py create mode 100644 vllm_omni/diffusion/cache/stepcache/state.py create mode 100644 vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py diff --git a/examples/offline_inference/dreamzero/client_schedule.py b/examples/offline_inference/dreamzero/client_schedule.py new file mode 100644 index 000000000..e4dbad55a --- /dev/null +++ b/examples/offline_inference/dreamzero/client_schedule.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""DreamZero AR client frame schedule aligned with upstream ``test_client_AR.py``.""" + +from __future__ import annotations + +import logging +from typing import Any + +import numpy as np + +# Frame schedule constants (matching upstream debug_inference.py / test_client_AR.py). +ACTION_HORIZON = 24 +DEFAULT_ACTION_DIM = 8 +RELATIVE_OFFSETS = [-23, -16, -8, 0] +DEFAULT_NUM_AR_CHUNKS = 15 + +CAMERA_FILES = { + "observation/exterior_image_0_left": "exterior_image_1_left.mp4", + "observation/exterior_image_1_left": "exterior_image_2_left.mp4", + "observation/wrist_image_left": "wrist_image_left.mp4", +} + + +def build_frame_schedule( + total_frames: int, + num_chunks: int, + *, + logger: logging.Logger | None = None, +) -> list[list[int]]: + """Build the frame index schedule for multi-frame chunks. + + Args: + total_frames: Number of frames available in each camera video. + num_chunks: Number of 4-frame chunks to schedule after the initial frame. + + Returns: + A list of frame-index lists. Each inner list has four indices. The + returned list may be shorter than ``num_chunks`` when the videos run + out of frames (upstream stops early instead of erroring). + """ + chunks: list[list[int]] = [] + current_frame = 23 + for _ in range(num_chunks): + indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS] + if indices[-1] >= total_frames: + if logger is not None: + logger.info( + "Frame %s >= %s, stopping at %s chunks", + indices[-1], + total_frames, + len(chunks), + ) + break + chunks.append(indices) + current_frame += ACTION_HORIZON + return chunks + + +def make_obs_from_video( + camera_frames: dict[str, np.ndarray], + frame_indices: list[int], + *, + prompt: str, + session_id: str, +) -> dict[str, Any]: + """Build an observation dict from real video frames. + + For one frame each image key is ``(H, W, 3)``. For four frames each key is + ``(4, H, W, 3)``. + """ + obs: dict[str, Any] = {} + for camera_key, all_frames in camera_frames.items(): + selected = all_frames[frame_indices] + obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected + + obs["observation/joint_position"] = np.zeros(7, dtype=np.float32) + obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32) + obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32) + obs["prompt"] = prompt + obs["session_id"] = session_id + return obs + + +def build_ar_observations( + camera_frames: dict[str, np.ndarray], + *, + prompt: str, + session_id: str, + num_chunks: int = DEFAULT_NUM_AR_CHUNKS, + repeat_chunk_observations: bool = False, + logger: logging.Logger | None = None, +) -> list[dict[str, Any]]: + """Build the AR observation sequence used by upstream ``test_client_AR.py``. + + Step 0 sends frame ``[0]``. Each subsequent step sends one 4-frame chunk. + + ``num_chunks`` counts only the 4-frame chunks after the initial frame, not + the total number of inferences. + """ + if num_chunks < 0: + raise ValueError("num_chunks must be non-negative") + + total_frames = min(frames.shape[0] for frames in camera_frames.values()) + observations = [ + make_obs_from_video( + camera_frames, + [0], + prompt=prompt, + session_id=session_id, + ) + ] + + chunk_schedule = build_frame_schedule(total_frames, num_chunks, logger=logger) + if repeat_chunk_observations and chunk_schedule and len(chunk_schedule) < num_chunks: + while len(chunk_schedule) < num_chunks: + chunk_schedule.append(chunk_schedule[-1]) + + for frame_indices in chunk_schedule: + observations.append( + make_obs_from_video( + camera_frames, + frame_indices, + prompt=prompt, + session_id=session_id, + ) + ) + return observations diff --git a/examples/offline_inference/dreamzero/export_prediction_video.py b/examples/offline_inference/dreamzero/export_prediction_video.py index 82f56ff26..4c560418c 100644 --- a/examples/offline_inference/dreamzero/export_prediction_video.py +++ b/examples/offline_inference/dreamzero/export_prediction_video.py @@ -11,6 +11,11 @@ from pathlib import Path import cv2 import numpy as np import torch +from client_schedule import ( + CAMERA_FILES, + DEFAULT_NUM_AR_CHUNKS, + build_ar_observations, +) from PIL import Image from vllm_omni import Omni @@ -20,13 +25,6 @@ from vllm_omni.outputs import OmniRequestOutput WORKER_EXTENSION = "vllm_omni.diffusion.models.dreamzero.video_export_worker.DreamZeroVideoExportWorkerExtension" ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets" DEFAULT_SESSION_PREFIX = "dreamzero-export" -RELATIVE_OFFSETS = [-23, -16, -8, 0] -ACTION_HORIZON = 24 -CAMERA_FILES = { - "observation/exterior_image_0_left": "exterior_image_1_left.mp4", - "observation/exterior_image_1_left": "exterior_image_2_left.mp4", - "observation/wrist_image_left": "wrist_image_left.mp4", -} def _parse_args() -> argparse.Namespace: @@ -49,6 +47,17 @@ def _parse_args() -> argparse.Namespace: default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan", ) parser.add_argument("--session-id", default=None) + parser.add_argument( + "--num-chunks", + type=int, + default=DEFAULT_NUM_AR_CHUNKS, + help="Number of 4-frame chunks to send after the initial frame (default: 15).", + ) + parser.add_argument( + "--repeat-chunk-observations", + action="store_true", + help="Repeat the last valid chunk observation when bundled assets run out of frames.", + ) parser.add_argument("--save-input-video", action="store_true") parser.add_argument("--save-gif", action="store_true") parser.add_argument("--save-actions", action="store_true") @@ -84,57 +93,25 @@ def _load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]: return camera_frames -def _build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]: - chunks: list[list[int]] = [] - current_frame = 23 - for _ in range(num_chunks): - indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS] - if indices[-1] >= total_frames: - break - chunks.append(indices) - current_frame += ACTION_HORIZON - return chunks - - -def _make_obs_from_video( - camera_frames: dict[str, np.ndarray], - frame_indices: list[int], - *, +def _build_observations( + video_dir: Path, prompt: str, session_id: str, -) -> dict: - obs: dict = {} - for camera_key, all_frames in camera_frames.items(): - selected = all_frames[frame_indices] - obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected - - obs["observation/joint_position"] = np.zeros(7, dtype=np.float32) - obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32) - obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32) - obs["prompt"] = prompt - obs["session_id"] = session_id - return obs - - -def _build_observations(video_dir: Path, prompt: str, session_id: str) -> tuple[dict[str, np.ndarray], list[dict]]: + *, + num_chunks: int, + repeat_chunk_observations: bool, +) -> tuple[dict[str, np.ndarray], list[dict]]: camera_frames = _load_camera_frames(video_dir) - total_frames = min(frames.shape[0] for frames in camera_frames.values()) - chunks = _build_frame_schedule(total_frames, 1) - observations = [ - _make_obs_from_video(camera_frames, [0], prompt=prompt, session_id=session_id), - ] - if chunks: - observations.append( - _make_obs_from_video( - camera_frames, - chunks[0], - prompt=prompt, - session_id=session_id, - ) - ) + observations = build_ar_observations( + camera_frames, + prompt=prompt, + session_id=session_id, + num_chunks=num_chunks, + repeat_chunk_observations=repeat_chunk_observations, + ) if len(observations) < 2: - raise RuntimeError("Need at least two DreamZero example observations to export a prediction video.") - return camera_frames, observations[:2] + raise RuntimeError("Need at least two DreamZero observations to export a prediction video.") + return camera_frames, observations def _extract_latents(output: OmniRequestOutput) -> torch.Tensor: @@ -229,6 +206,26 @@ def _run_generation( return omni, outputs +def _decode_accumulated_with_worker(omni: Omni, session_id: str) -> np.ndarray: + """Decode server-side accumulated latents when the pipeline supports it.""" + stage_client = omni.engine.stage_clients[0] + engine = getattr(stage_client, "_engine", None) + if engine is None: + raise RuntimeError("DreamZero export requires inline diffusion stage access.") + + decoded = engine.executor.collective_rpc( + "decode_accumulated_video_latents_to_uint8", + args=(session_id,), + unique_reply_rank=0, + exec_all_ranks=True, + ) + if isinstance(decoded, torch.Tensor): + decoded = decoded.numpy() + if not isinstance(decoded, np.ndarray): + raise TypeError(f"Unexpected decoded output type: {type(decoded)!r}") + return decoded + + def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray: stage_client = omni.engine.stage_clients[0] engine = getattr(stage_client, "_engine", None) @@ -248,6 +245,15 @@ def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray: return decoded +def _decode_video(omni: Omni, session_id: str, outputs: list[OmniRequestOutput]) -> np.ndarray: + """Decode full AR rollout by concatenating per-step latents from all outputs.""" + latent_steps = [_extract_latents(output) for output in outputs] + latent_t = [step.shape[2] for step in latent_steps] + print(f"Decoding {len(latent_steps)} latent steps, latent_t={latent_t}, total_t={sum(latent_t)}") + full_latents = torch.cat(latent_steps, dim=2) + return _decode_with_worker(omni, full_latents) + + def main() -> None: args = _parse_args() session_id = args.session_id or f"{DEFAULT_SESSION_PREFIX}-{uuid.uuid4()}" @@ -256,6 +262,12 @@ def main() -> None: video_dir=args.video_dir, prompt=args.prompt, session_id=session_id, + num_chunks=args.num_chunks, + repeat_chunk_observations=args.repeat_chunk_observations, + ) + print( + f"Running {len(observations)} inferences " + f"(1 initial + {len(observations) - 1} chunks; requested {args.num_chunks} post-initial chunks)" ) args.output_dir.mkdir(parents=True, exist_ok=True) @@ -273,17 +285,19 @@ def main() -> None: deploy_config_path=args.deploy_config, observations=observations, ) - latent_steps = [_extract_latents(output) for output in outputs] - full_latents = torch.cat(latent_steps, dim=2) - frames = _decode_with_worker(omni, full_latents) + frames = _decode_video(omni, session_id, outputs) finally: if omni is not None: omni.close() + expected_frames = 8 * len(observations) + 1 mp4_path = args.output_dir / f"{args.output_stem}.mp4" _write_mp4(mp4_path, frames, fps=args.fps) - print(f"SAVED_MP4={mp4_path}") + print( + f"SAVED_MP4={mp4_path} frames={len(frames)} " + f"expected≈{expected_frames} duration≈{len(frames) / max(args.fps, 1):.1f}s" + ) if args.save_gif: gif_path = args.output_dir / f"{args.output_stem}.gif" @@ -294,8 +308,10 @@ def main() -> None: npz_path = args.output_dir / f"{args.output_stem}_actions.npz" np.savez( npz_path, - step0=np.asarray(outputs[0].multimodal_output.get("actions")), - step1=np.asarray(outputs[1].multimodal_output.get("actions")), + **{ + f"step{index}": np.asarray(output.multimodal_output.get("actions")) + for index, output in enumerate(outputs) + }, ) print(f"SAVED_ACTIONS={npz_path}") diff --git a/tests/diffusion/cache/test_stepcache.py b/tests/diffusion/cache/test_stepcache.py new file mode 100644 index 000000000..1fc3bbb59 --- /dev/null +++ b/tests/diffusion/cache/test_stepcache.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for StepCacheBackend and StepCacheState.""" + +from unittest.mock import Mock + +import pytest +import torch + +from vllm_omni.diffusion.cache.selector import get_cache_backend +from vllm_omni.diffusion.cache.stepcache import ( + STEP_CACHE_CONFIG_ATTR, + STEP_CACHE_STATE_ATTR, + StepCacheBackend, + StepCacheConfig, + StepCacheState, + is_stepcache_active, +) +from vllm_omni.diffusion.data import DiffusionCacheConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_prev_predictions(sim: float, dim: int = 8) -> list[tuple[torch.Tensor]]: + v_prev = torch.randn(1, dim) + if sim >= 0.99: + v_last = v_prev.clone() + else: + v_last = torch.randn(1, dim) + return [(v_prev,), (v_last,)] + + +class TestStepCacheConfig: + def test_from_diffusion_cache_config_respects_step_cache_dit_enabled(self): + enabled = StepCacheConfig.from_diffusion_cache_config(DiffusionCacheConfig(step_cache_dit_enabled=True)) + disabled = StepCacheConfig.from_diffusion_cache_config(DiffusionCacheConfig(step_cache_dit_enabled=False)) + assert enabled.enabled is True + assert disabled.enabled is False + + +class TestStepCacheState: + def test_disabled_always_runs(self): + state = StepCacheState(StepCacheConfig(enabled=False)) + assert state.should_run_step(_make_prev_predictions(0.99)) is True + + def test_insufficient_history_always_runs(self): + state = StepCacheState(StepCacheConfig(min_history_steps=2)) + assert state.should_run_step([(torch.randn(1, 4),)]) is True + + def test_high_similarity_skips_with_countdown(self): + state = StepCacheState( + StepCacheConfig( + sim_thresholds=(0.95,), + skip_countdowns=(2,), + ) + ) + preds = _make_prev_predictions(0.99) + assert state.should_run_step(preds) is False + assert state.skip_countdown == 2 + assert state.should_run_step(preds) is False + assert state.should_run_step(preds) is True + + def test_low_similarity_runs(self): + state = StepCacheState(StepCacheConfig()) + preds = _make_prev_predictions(0.5) + assert state.should_run_step(preds) is True + + def test_trim_history(self): + state = StepCacheState(StepCacheConfig(max_history=2)) + history: list[tuple[torch.Tensor]] = [(torch.randn(1, 2),) for _ in range(4)] + state.trim_history(history) + assert len(history) == 2 + + +class TestStepCacheBackend: + def test_enable_dreamzero_pipeline(self): + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "DreamZeroPipeline" + + backend = StepCacheBackend(DiffusionCacheConfig()) + backend.enable(mock_pipeline) + + assert backend.enabled is True + cfg = getattr(mock_pipeline, STEP_CACHE_CONFIG_ATTR) + assert isinstance(cfg, StepCacheConfig) + assert is_stepcache_active(mock_pipeline) + state = getattr(mock_pipeline, STEP_CACHE_STATE_ATTR) + assert isinstance(state, StepCacheState) + + def test_refresh_resets_state(self): + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "DreamZeroPipeline" + backend = StepCacheBackend(DiffusionCacheConfig()) + backend.enable(mock_pipeline) + state = getattr(mock_pipeline, STEP_CACHE_STATE_ATTR) + state.skip_countdown = 3 + backend.refresh(mock_pipeline, num_inference_steps=16, verbose=False) + assert state.skip_countdown == 0 + + def test_enable_unsupported_pipeline_raises(self): + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "QwenImagePipeline" + backend = StepCacheBackend(DiffusionCacheConfig()) + + with pytest.raises(ValueError, match="step_cache backend does not support"): + backend.enable(mock_pipeline) + + def test_selector_returns_backend(self): + backend = get_cache_backend("step_cache", {"velocity_sim_thresholds": [0.9]}) + assert isinstance(backend, StepCacheBackend) + assert backend.config.velocity_sim_thresholds == [0.9] + + def test_legacy_backend_aliases(self): + assert isinstance(get_cache_backend("stepcache", {}), StepCacheBackend) + assert isinstance(get_cache_backend("step_cache_dit", {}), StepCacheBackend) diff --git a/tests/e2e/online_serving/test_dreamzero_expansion.py b/tests/e2e/online_serving/test_dreamzero_expansion.py index bc3215891..79429c123 100644 --- a/tests/e2e/online_serving/test_dreamzero_expansion.py +++ b/tests/e2e/online_serving/test_dreamzero_expansion.py @@ -28,8 +28,6 @@ test_params = [ server_args=[ "--deploy-config", "vllm_omni/deploy/dreamzero_tp1_cfg2.yaml", - "--enforce-eager", - "--disable-log-stats", ], env_dict={ "ATTENTION_BACKEND": "torch", diff --git a/tests/entrypoints/openai_api/test_openpi_connection.py b/tests/entrypoints/openai_api/test_openpi_connection.py index 83d17a28c..398ae6494 100644 --- a/tests/entrypoints/openai_api/test_openpi_connection.py +++ b/tests/entrypoints/openai_api/test_openpi_connection.py @@ -88,6 +88,47 @@ def test_unpack_accepts_msgpack_numpy_marker_dicts(): np.testing.assert_allclose(decoded[b"actions"], np.asarray([[1.0, 0.0]], dtype=np.float32)) +def test_unpack_accepts_openpi_client_ndarray_markers(): + image = np.arange(6, dtype=np.uint8).reshape(2, 3) + payload = { + "observation/exterior_image_0_left": { + b"__ndarray__": True, + b"dtype": image.dtype.str, + b"shape": image.shape, + b"data": image.tobytes(), + } + } + + decoded = openpi_connection._unpack_numpy(payload) + + np.testing.assert_array_equal(decoded["observation/exterior_image_0_left"], image) + assert decoded["observation/exterior_image_0_left"].dtype == np.uint8 + + +def test_unpack_openpi_client_packed_observation(): + openpi_msgpack = pytest.importorskip("openpi_client.msgpack_numpy") + image = np.zeros((180, 320, 3), dtype=np.uint8) + obs = { + "observation/exterior_image_0_left": image, + "observation/joint_position": np.zeros(7, dtype=np.float32), + } + + decoded = openpi_connection._unpack(openpi_msgpack.packb(obs)) + + np.testing.assert_array_equal(decoded["observation/exterior_image_0_left"], image) + np.testing.assert_allclose(decoded["observation/joint_position"], obs["observation/joint_position"]) + + +def test_pack_uses_openpi_client_ndarray_markers(): + openpi_msgpack = pytest.importorskip("openpi_client.msgpack_numpy") + actions = np.asarray([[1.0, 2.0]], dtype=np.float32) + + packed = openpi_connection._pack(actions) + decoded = openpi_msgpack.unpackb(packed) + + np.testing.assert_allclose(decoded, actions) + + def test_unpack_leaves_user_dict_without_numpy_kind_marker_unchanged(): payload = { "metadata": { diff --git a/vllm_omni/deploy/dreamzero.yaml b/vllm_omni/deploy/dreamzero.yaml index e77afad84..079101ba0 100644 --- a/vllm_omni/deploy/dreamzero.yaml +++ b/vllm_omni/deploy/dreamzero.yaml @@ -12,7 +12,7 @@ stages: - stage_id: 0 devices: "0" max_num_seqs: 1 - enforce_eager: true + enforce_eager: false model_class_name: DreamZeroPipeline model_config: default_robot_embodiment: roboarena @@ -23,3 +23,8 @@ stages: needs_stereo_camera: false needs_session_id: true action_space: joint_position + cache_backend: step_cache + cache_config: + step_cache_dit_enabled: true + velocity_sim_thresholds: [0.95, 0.93] + velocity_skip_countdowns: [4, 2] diff --git a/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml b/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml index ed9988c92..a453d43f3 100644 --- a/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml +++ b/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml @@ -9,7 +9,7 @@ stages: - stage_id: 0 devices: "0,1" max_num_seqs: 1 - enforce_eager: true + enforce_eager: false model_class_name: DreamZeroPipeline parallel_config: tensor_parallel_size: 1 @@ -23,3 +23,8 @@ stages: needs_stereo_camera: false needs_session_id: true action_space: joint_position + cache_backend: step_cache + cache_config: + step_cache_dit_enabled: true + velocity_sim_thresholds: [0.95, 0.93] + velocity_skip_countdowns: [4, 2] diff --git a/vllm_omni/diffusion/cache/base.py b/vllm_omni/diffusion/cache/base.py index d78076224..aafbf0ba7 100644 --- a/vllm_omni/diffusion/cache/base.py +++ b/vllm_omni/diffusion/cache/base.py @@ -13,6 +13,8 @@ Main cache backend implementations: the cache-dit library. Inherits from CacheBackend. Used via cache_backend="cache_dit". 2. TeaCacheBackend: Hook-based backend for TeaCache acceleration. Inherits from CacheBackend. Used via cache_backend="tea_cache". +3. StepCacheBackend: Velocity cosine step-skipping for DreamZero. Inherits from + CacheBackend. Used via cache_backend="step_cache". All backends implement the same interface: - enable(pipeline): Enable cache on the pipeline diff --git a/vllm_omni/diffusion/cache/selector.py b/vllm_omni/diffusion/cache/selector.py index e6c2e70e3..3b002abe0 100644 --- a/vllm_omni/diffusion/cache/selector.py +++ b/vllm_omni/diffusion/cache/selector.py @@ -3,6 +3,7 @@ from typing import Any from vllm_omni.diffusion.cache.base import CacheBackend from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend +from vllm_omni.diffusion.cache.stepcache import StepCacheBackend from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend from vllm_omni.diffusion.data import DiffusionCacheConfig @@ -14,14 +15,15 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack - cache_dit: Uses CacheDiTBackend with enable()/refresh() interface - tea_cache: Uses TeaCacheBackend with enable()/refresh() interface - mag_cache: Uses MagCacheBackend with enable()/refresh() interface + - step_cache: Uses StepCacheBackend for DreamZero velocity step skip Args: - cache_backend: Cache backend name ("cache_dit", "tea_cache", "mag_cache", or None). + cache_backend: Cache backend name ("cache_dit", "tea_cache", + "mag_cache", "step_cache", or None). cache_config: Cache configuration (dict or DiffusionCacheConfig instance). Returns: - Cache backend instance (CacheDiTBackend, TeaCacheBackend, or MagCacheBackend) - if cache_backend is set, None otherwise. + Cache backend instance if cache_backend is set, None otherwise. Raises: ValueError: If cache_backend is unsupported. @@ -38,7 +40,10 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack return TeaCacheBackend(cache_config) elif cache_backend == "mag_cache": return MagCacheBackend(cache_config) + elif cache_backend in ("step_cache", "stepcache", "step_cache_dit"): + return StepCacheBackend(cache_config) else: raise ValueError( - f"Unsupported cache backend: {cache_backend}. Supported: 'cache_dit', 'tea_cache', 'mag_cache'" + f"Unsupported cache backend: {cache_backend}. " + "Supported: 'cache_dit', 'tea_cache', 'mag_cache', 'step_cache'" ) diff --git a/vllm_omni/diffusion/cache/stepcache/__init__.py b/vllm_omni/diffusion/cache/stepcache/__init__.py new file mode 100644 index 000000000..a7e1852b0 --- /dev/null +++ b/vllm_omni/diffusion/cache/stepcache/__init__.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +StepCache: velocity-based step skipping for DreamZero-style DiT denoising. + +Unlike block-level Cache-DiT, this backend skips entire DiT forwards during the +pipeline denoise loop when successive velocity predictions are highly aligned. + +Usage: + from vllm_omni import Omni + + omni = Omni( + model="...", + cache_backend="step_cache", + ) + + # Alternative: environment variable + # export DIFFUSION_CACHE_BACKEND=step_cache +""" + +from vllm_omni.diffusion.cache.stepcache.backend import ( + CUSTOM_STEP_CACHE_DIT_ENABLERS, + CUSTOM_STEPCACHE_ENABLERS, + STEP_CACHE_CONFIG_ATTR, + STEP_CACHE_DIT_CONFIG_ATTR, + STEP_CACHE_DIT_STATE_ATTR, + STEP_CACHE_STATE_ATTR, + StepCacheBackend, + StepCacheDiTBackend, + StepCacheDiTConfig, + StepCacheDiTState, + enable_dreamzero_step_cache_dit, + enable_dreamzero_stepcache, + get_step_cache_dit_config, + get_step_cache_dit_state, + get_stepcache_config, + get_stepcache_state, + is_step_cache_dit_active, + is_stepcache_active, +) +from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig +from vllm_omni.diffusion.cache.stepcache.state import StepCacheState + +__all__ = [ + "CUSTOM_STEP_CACHE_DIT_ENABLERS", + "CUSTOM_STEPCACHE_ENABLERS", + "STEP_CACHE_CONFIG_ATTR", + "STEP_CACHE_DIT_CONFIG_ATTR", + "STEP_CACHE_DIT_STATE_ATTR", + "STEP_CACHE_STATE_ATTR", + "StepCacheBackend", + "StepCacheConfig", + "StepCacheDiTBackend", + "StepCacheDiTConfig", + "StepCacheDiTState", + "StepCacheState", + "enable_dreamzero_step_cache_dit", + "enable_dreamzero_stepcache", + "get_step_cache_dit_config", + "get_step_cache_dit_state", + "get_stepcache_config", + "get_stepcache_state", + "is_step_cache_dit_active", + "is_stepcache_active", +] diff --git a/vllm_omni/diffusion/cache/stepcache/backend.py b/vllm_omni/diffusion/cache/stepcache/backend.py new file mode 100644 index 000000000..517d2369a --- /dev/null +++ b/vllm_omni/diffusion/cache/stepcache/backend.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +StepCache backend implementation. + +This module provides the stepcache backend that implements the CacheBackend +interface for DreamZero-style velocity-based step skipping in the denoise loop. +""" + +from __future__ import annotations + +from typing import Any + +from vllm.logger import init_logger + +from vllm_omni.diffusion.cache.base import CacheBackend +from vllm_omni.diffusion.cache.stepcache.config import ( + STEP_CACHE_CONFIG_ATTR, + STEP_CACHE_STATE_ATTR, + StepCacheConfig, +) +from vllm_omni.diffusion.cache.stepcache.state import StepCacheState +from vllm_omni.diffusion.data import DiffusionCacheConfig + +logger = init_logger(__name__) + +# Backward-compatible aliases for older imports and pipeline attribute names. +STEP_CACHE_DIT_CONFIG_ATTR = STEP_CACHE_CONFIG_ATTR +STEP_CACHE_DIT_STATE_ATTR = STEP_CACHE_STATE_ATTR +StepCacheDiTConfig = StepCacheConfig +StepCacheDiTState = StepCacheState + + +def is_stepcache_active(pipeline: Any) -> bool: + """Return True when stepcache is enabled on *pipeline*.""" + cfg = getattr(pipeline, STEP_CACHE_CONFIG_ATTR, None) + return isinstance(cfg, StepCacheConfig) and cfg.enabled + + +def get_stepcache_config(pipeline: Any) -> StepCacheConfig | None: + cfg = getattr(pipeline, STEP_CACHE_CONFIG_ATTR, None) + return cfg if isinstance(cfg, StepCacheConfig) else None + + +def get_stepcache_state(pipeline: Any) -> StepCacheState | None: + state = getattr(pipeline, STEP_CACHE_STATE_ATTR, None) + return state if isinstance(state, StepCacheState) else None + + +def is_step_cache_dit_active(pipeline: Any) -> bool: + return is_stepcache_active(pipeline) + + +def get_step_cache_dit_config(pipeline: Any) -> StepCacheConfig | None: + return get_stepcache_config(pipeline) + + +def get_step_cache_dit_state(pipeline: Any) -> StepCacheState | None: + return get_stepcache_state(pipeline) + + +def enable_dreamzero_stepcache(pipeline: Any, config: DiffusionCacheConfig) -> None: + """Enable stepcache for DreamZeroPipeline.""" + cache_config = StepCacheConfig.from_diffusion_cache_config(config) + setattr(pipeline, STEP_CACHE_CONFIG_ATTR, cache_config) + setattr(pipeline, STEP_CACHE_STATE_ATTR, StepCacheState(cache_config)) + + logger.info( + "stepcache enabled on DreamZeroPipeline (enabled=%s, thresholds=%s, countdowns=%s)", + cache_config.enabled, + cache_config.sim_thresholds, + cache_config.skip_countdowns, + ) + + +enable_dreamzero_step_cache_dit = enable_dreamzero_stepcache + +CUSTOM_STEPCACHE_ENABLERS = { + "DreamZeroPipeline": enable_dreamzero_stepcache, +} + +CUSTOM_STEP_CACHE_DIT_ENABLERS = CUSTOM_STEPCACHE_ENABLERS + + +class StepCacheBackend(CacheBackend): + """ + Velocity cosine step-skipping cache backend for DreamZero. + + Attaches :class:`StepCacheConfig` and :class:`StepCacheState` to supported + pipelines. The denoise loop calls :meth:`StepCacheState.should_run_step` to + decide whether to run ``predict_noise``. + + Example: + >>> from vllm_omni.diffusion.data import DiffusionCacheConfig + >>> backend = StepCacheBackend(DiffusionCacheConfig()) + >>> backend.enable(pipeline) + >>> backend.refresh(pipeline, num_inference_steps=16) + """ + + def enable(self, pipeline: Any) -> None: + """Enable stepcache on a supported pipeline.""" + pipeline_type = pipeline.__class__.__name__ + + if pipeline_type in CUSTOM_STEPCACHE_ENABLERS: + logger.info("Using custom stepcache enabler for model: %s", pipeline_type) + CUSTOM_STEPCACHE_ENABLERS[pipeline_type](pipeline, self.config) + else: + raise ValueError( + f"step_cache backend does not support {pipeline_type}. Supported: {sorted(CUSTOM_STEPCACHE_ENABLERS)}" + ) + + self.enabled = True + + def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + """Refresh stepcache state for a new generation.""" + state = get_stepcache_state(pipeline) + if state is not None: + state.reset() + if verbose: + logger.debug( + "stepcache state refreshed (num_inference_steps=%d)", + num_inference_steps, + ) + elif verbose and is_stepcache_active(pipeline): + logger.warning("stepcache config active but state not found on pipeline") + + +StepCacheDiTBackend = StepCacheBackend diff --git a/vllm_omni/diffusion/cache/stepcache/config.py b/vllm_omni/diffusion/cache/stepcache/config.py new file mode 100644 index 000000000..00a7984ae --- /dev/null +++ b/vllm_omni/diffusion/cache/stepcache/config.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Configuration for step-level DiT velocity caching.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from vllm_omni.diffusion.data import DiffusionCacheConfig + +# Pipeline attributes set by StepCacheBackend.enable(). +STEP_CACHE_CONFIG_ATTR = "_stepcache_config" +STEP_CACHE_STATE_ATTR = "_stepcache_state" + + +@dataclass(frozen=True) +class StepCacheConfig: + """Runtime config for velocity-based step skipping. + + Skips entire DiT forwards when successive video velocity predictions are + highly aligned (cosine similarity above configured thresholds). + + Reference: DreamZero paper DiT Caching / dreamzero.git ``should_run_model``. + """ + + enabled: bool = True + min_history_steps: int = 2 + max_history: int = 2 + sim_thresholds: tuple[float, ...] = (0.95, 0.93) + skip_countdowns: tuple[int, ...] = (4, 2) + + @classmethod + def from_diffusion_cache_config(cls, config: DiffusionCacheConfig) -> StepCacheConfig: + enabled = config.step_cache_dit_enabled + + thresholds = tuple(config.velocity_sim_thresholds) + countdowns = tuple(config.velocity_skip_countdowns) + if len(thresholds) != len(countdowns): + raise ValueError( + "velocity_sim_thresholds and velocity_skip_countdowns must have the same length; " + f"got {len(thresholds)} and {len(countdowns)}." + ) + + return cls( + enabled=enabled, + min_history_steps=config.step_cache_dit_min_history, + max_history=config.step_cache_dit_max_history, + sim_thresholds=thresholds, + skip_countdowns=countdowns, + ) diff --git a/vllm_omni/diffusion/cache/stepcache/state.py b/vllm_omni/diffusion/cache/stepcache/state.py new file mode 100644 index 000000000..db786d4f2 --- /dev/null +++ b/vllm_omni/diffusion/cache/stepcache/state.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Mutable state for step-level DiT velocity caching.""" + +from __future__ import annotations + +import torch + +from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig + + +class StepCacheState: + """Per-generation mutable state (skip countdown).""" + + def __init__(self, config: StepCacheConfig) -> None: + self.config = config + self.skip_countdown = 0 + + def reset(self) -> None: + self.skip_countdown = 0 + + def should_run_step(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> bool: + """Return True when the DiT forward should execute at this scheduler step.""" + if not self.config.enabled: + return True + + if len(prev_predictions) < self.config.min_history_steps: + return True + + if self.skip_countdown > 1: + self.skip_countdown -= 1 + return False + if self.skip_countdown == 1: + self.skip_countdown = 0 + return True + + v_last = prev_predictions[-1][0].flatten(1).float() + v_prev = prev_predictions[-2][0].flatten(1).float() + sim = torch.nn.functional.cosine_similarity(v_last, v_prev, dim=1).mean() + + for threshold, countdown in zip(self.config.sim_thresholds, self.config.skip_countdowns): + if sim > threshold: + self.skip_countdown = countdown + return False + + return True + + def trim_history(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> None: + while len(prev_predictions) > self.config.max_history: + prev_predictions.pop(0) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index c4cca97bb..bf294fdfe 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -302,6 +302,8 @@ class DiffusionCacheConfig: scm_steps_mask_policy, scm_steps_policy - MagCache: mag_threshold, mag_max_skip_steps, mag_retention_ratio, mag_ratios, mag_calibrate + - step_cache: step_cache_dit_enabled, velocity_sim_thresholds, + velocity_skip_countdowns, step_cache_dit_min_history Example: >>> # From dict (user-facing API) - partial config uses defaults for missing keys @@ -375,6 +377,13 @@ class DiffusionCacheConfig: # "repeat" refreshes every force_refresh_step_hint steps. force_refresh_step_policy: str = "once" + # step_cache parameters [step_cache only] — DreamZero velocity schedule + step_cache_dit_enabled: bool = True + velocity_sim_thresholds: list[float] = field(default_factory=lambda: [0.95, 0.93]) + velocity_skip_countdowns: list[int] = field(default_factory=lambda: [4, 2]) + step_cache_dit_min_history: int = 2 + step_cache_dit_max_history: int = 2 + # Additional parameters that may be passed but not explicitly defined _extra_params: dict[str, Any] = field(default_factory=dict, repr=False) diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py index 1d1429534..bbff12373 100644 --- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py +++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py @@ -23,6 +23,10 @@ from huggingface_hub import hf_hub_download from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm_omni.diffusion.cache.stepcache import ( + get_stepcache_state, + is_stepcache_active, +) from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( DistributedAutoencoderKLWan, @@ -55,10 +59,6 @@ logger = logging.getLogger(__name__) MAX_DREAMZERO_SESSIONS = 64 -# --------------------------------------------------------------------------- -# --------------------------------------------------------------------------- - - class VideoActionScheduler: """Wraps video + action schedulers into single .step() interface.""" @@ -205,6 +205,9 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): self._max_session_states = MAX_DREAMZERO_SESSIONS self.state = self._get_or_create_state("default") + # DiT step cache is configured by StepCacheBackend + # (cache_backend="step_cache") via pipeline._stepcache_config. + # Keep runtime inference settings separate from the training-time config. self.num_inference_steps: int = model_config.get( "num_inference_steps", @@ -306,6 +309,27 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): def predict_noise(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]: """Call CausalWanModel, return (video_pred, action_pred).""" + video_pred, action_pred = self._predict_noise_eager(kwargs) + + if is_stepcache_active(self): + video_pred = video_pred.clone() + if action_pred is not None: + action_pred = action_pred.clone() + + if action_pred is None: + batch_size = kwargs["hidden_states"].shape[0] + action_pred = torch.empty( + batch_size, + 0, + self.transformer.action_dim, + device=video_pred.device, + dtype=video_pred.dtype, + ) + return (video_pred, action_pred) + + def _predict_noise_eager(self, kwargs: dict) -> tuple[torch.Tensor, torch.Tensor | None]: + """Eager DiT forward; also handles KV-cache write-back on prefill.""" + self._cudagraph_mark_step_begin() video_pred, action_pred, updated_kv_caches = self.transformer( x=kwargs["hidden_states"], timestep=kwargs["timestep_video"], @@ -327,19 +351,118 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): for i, kv in enumerate(updated_kv_caches): state.update_kv_cache(i, kv, is_negative=is_neg) - video_pred = video_pred.clone() - if action_pred is not None: - action_pred = action_pred.clone() - else: - batch_size = kwargs["hidden_states"].shape[0] - action_pred = torch.empty( - batch_size, - 0, - self.transformer.action_dim, - device=video_pred.device, - dtype=video_pred.dtype, - ) # CFG-parallel-safe dummy action pred - return (video_pred, action_pred) + return video_pred, action_pred + + # ----------------------------------------------------------------------- + # torch.compile setup (paper D.2 extended: encoders + VAE + DiT blocks) + # ----------------------------------------------------------------------- + + def setup_compile(self) -> None: + """Compile DreamZero encoders, VAE, and per-block DiT for inference. + + Paper D.2 uses ``mode=reduce-overhead``, ``fullgraph=True``, ``dynamic=False`` + on text/image/VAE and DiT. VAE decode uses a tensor feat_cache patch and + compiles ``decoder.forward`` (not ``_decode``, which has a Python frame loop). + DiT blocks use per-block ``fullgraph=True``. + """ + if not torch.cuda.is_available(): + logger.info("DreamZero setup_compile skipped: CUDA not available.") + return + + from vllm_omni.diffusion.models.dreamzero.wan_vae_feat_cache_patch import ( + apply_wan_vae_feat_cache_tensor_patch, + ) + + apply_wan_vae_feat_cache_tensor_patch() + + compile_ro = {"mode": "reduce-overhead", "fullgraph": True, "dynamic": False} + # DiT blocks: default avoids CUDAGraph overwrite on modulation tensors; encoders use reduce-overhead. + dit_compile = {"mode": "default", "fullgraph": True, "dynamic": False} + + logger.info( + "DreamZero: torch.compile text/image/VAE encode + per-block DiT (encoders reduce-overhead, DiT default)." + ) + + try: + self.text_encoder.forward = torch.compile(self.text_encoder.forward, **compile_ro) + except Exception as exc: + logger.warning("DreamZero: text_encoder compile failed (%s); skipping.", exc) + + try: + self.image_encoder.model.visual.forward = torch.compile( + self.image_encoder.model.visual.forward, + **compile_ro, + ) + except Exception as exc: + logger.warning("DreamZero: image_encoder compile failed (%s); skipping.", exc) + + try: + self.vae.encode = torch.compile(self.vae.encode, **compile_ro) + except Exception as exc: + logger.warning("DreamZero: vae.encode compile failed (%s); skipping.", exc) + + compiled_blocks = 0 + for block in self.transformer.blocks: + try: + block.forward = torch.compile(block.forward, **dit_compile) + compiled_blocks += 1 + except Exception as exc: + logger.warning( + "DreamZero: transformer block %d compile failed (%s); leaving remaining eager.", + compiled_blocks, + exc, + ) + break + if compiled_blocks: + logger.info( + "DreamZero: compiled %d/%d transformer blocks.", + compiled_blocks, + len(self.transformer.blocks), + ) + + self.warmup_compile() + + def warmup_compile(self) -> None: + """Warm up compiled text/image/VAE paths before timed inference.""" + if not torch.cuda.is_available(): + return + + device = next(self.text_encoder.parameters()).device + with torch.inference_mode(): + try: + text_tokens = torch.zeros(1, 16, dtype=torch.long, device=device) + attention_mask = torch.ones_like(text_tokens) + self._encode_text(text_tokens, attention_mask) + except Exception as exc: + logger.warning("DreamZero compile warmup (text_encoder) skipped: %s", exc) + + try: + image = torch.zeros(1, 1, 3, 180, 320, dtype=torch.bfloat16, device=device) + self._encode_image(image, self.num_frames, 180, 320) + except Exception as exc: + logger.warning("DreamZero compile warmup (image_encoder) skipped: %s", exc) + + try: + latent_h, latent_w = 180 // 8, 320 // 8 + dummy_latent = torch.zeros( + 1, + 16, + self.num_frame_per_block, + latent_h * 2, + latent_w * 2, + dtype=self.vae.dtype, + device=device, + ) + mean = self.vae_latents_mean.to(device=device, dtype=self.vae.dtype) + inv_std = self.vae_latents_inv_std.to(device=device, dtype=self.vae.dtype) + dummy_latent_denorm = dummy_latent / inv_std + mean + self._cudagraph_mark_step_begin() + self.vae.decode(dummy_latent_denorm, return_dict=False) + except Exception as exc: + logger.warning("DreamZero compile warmup (vae decode) skipped: %s", exc) + + torch.accelerator.synchronize(device) + logger.info("DreamZero compile warmup finished (text / image / vae decode).") def combine_cfg_noise( self, @@ -388,12 +511,20 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): videos = videos.reshape(b, t, c, h, w).permute(0, 2, 1, 3, 4) return videos.to(dtype=torch.bfloat16) + @staticmethod + def _cudagraph_mark_step_begin() -> None: + try: + torch.compiler.cudagraph_mark_step_begin() + except Exception: + pass + # ----------------------------------------------------------------------- # Text encoding # ----------------------------------------------------------------------- def _encode_text(self, text_tokens: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: """Encode text prompt via UMT5.""" + self._cudagraph_mark_step_begin() seq_lens = attention_mask.gt(0).sum(dim=1).long() prompt_emb = self.text_encoder( text_tokens, @@ -422,6 +553,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): batch_size = image.shape[0] with torch.amp.autocast(dtype=torch.bfloat16, device_type=device.type): + self._cudagraph_mark_step_begin() clip_context = self.image_encoder.encode_image(image) msk = torch.ones(batch_size, num_frames, height // 8, width // 8, device=device) @@ -469,8 +601,23 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): inv_std = self.vae_latents_inv_std.to(device=vae_device, dtype=vae_dtype) latents = latents / inv_std + mean with torch.no_grad(): + self._cudagraph_mark_step_begin() return self.vae.decode(latents, return_dict=False)[0] + def decode_accumulated_video_latents(self, session_id: str | None = None) -> torch.Tensor: + """Decode all AR-chunk latents accumulated for ``session_id``.""" + state = self._get_or_create_state(session_id) + latents = state.get_concatenated_video_latents() + if latents is None: + session_key = str(session_id or "default") + raise RuntimeError(f"No accumulated video latents for session {session_key!r}.") + return self.decode_video_latents(latents) + + def clear_accumulated_video_latents(self, session_id: str | None = None) -> None: + """Clear accumulated video latents for ``session_id`` without resetting KV state.""" + state = self._get_or_create_state(session_id) + state.clear_video_latents() + # ----------------------------------------------------------------------- # KV cache prefill # ----------------------------------------------------------------------- @@ -635,6 +782,12 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): noisy_input = video_latents noisy_input_action = action_latents + # ---- step_cache (StepCacheBackend / dreamzero.git) ---- + _cached_flow_pred: torch.Tensor | None = None + _cached_flow_pred_action: torch.Tensor | None = None + _prev_predictions: list[tuple[torch.Tensor]] = [] + _step_cache = get_stepcache_state(self) if is_stepcache_active(self) else None + for index in range(len(timesteps_video)): video_timestep = timesteps_video[index] action_timestep = timesteps_action[index] @@ -663,46 +816,58 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): else: y = state.ys[:, :, -self.num_frame_per_block :] - positive_kwargs = dict( - hidden_states=noisy_input.transpose(1, 2), - timestep_video=timestep, - encoder_hidden_states=prompt_embeds, - kv_cache=state.get_kv_caches(False), - crossattn_cache=state.get_crossattn_caches(False), - y=y, - clip_feature=state.clip_feas, - action=noisy_input_action, - timestep_action=timestep_action, - is_negative=False, - **common_kwargs, - ) - - # Negative (uncond) kwargs - if do_true_cfg and negative_prompt_embeds is not None: - negative_kwargs = dict( + run_dit = _step_cache is None or _step_cache.should_run_step(_prev_predictions) + if run_dit: + positive_kwargs = dict( hidden_states=noisy_input.transpose(1, 2), timestep_video=timestep, - encoder_hidden_states=negative_prompt_embeds, - kv_cache=state.get_kv_caches(True), - crossattn_cache=state.get_crossattn_caches(True), + encoder_hidden_states=prompt_embeds, + kv_cache=state.get_kv_caches(False), + crossattn_cache=state.get_crossattn_caches(False), y=y, clip_feature=state.clip_feas, action=noisy_input_action, timestep_action=timestep_action, - is_negative=True, + is_negative=False, **common_kwargs, ) - else: - negative_kwargs = None - noise_pred = self.predict_noise_maybe_with_cfg( - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - do_true_cfg=do_true_cfg, - true_cfg_scale=self.cfg_scale, - cfg_normalize=False, - ) - flow_pred, flow_pred_action = noise_pred + if do_true_cfg and negative_prompt_embeds is not None: + negative_kwargs = dict( + hidden_states=noisy_input.transpose(1, 2), + timestep_video=timestep, + encoder_hidden_states=negative_prompt_embeds, + kv_cache=state.get_kv_caches(True), + crossattn_cache=state.get_crossattn_caches(True), + y=y, + clip_feature=state.clip_feas, + action=noisy_input_action, + timestep_action=timestep_action, + is_negative=True, + **common_kwargs, + ) + else: + negative_kwargs = None + + noise_pred = self.predict_noise_maybe_with_cfg( + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + do_true_cfg=do_true_cfg, + true_cfg_scale=self.cfg_scale, + cfg_normalize=False, + ) + flow_pred, flow_pred_action = noise_pred + _cached_flow_pred = flow_pred + _cached_flow_pred_action = flow_pred_action + + _prev_predictions.append((flow_pred,)) + if _step_cache is not None: + _step_cache.trim_history(_prev_predictions) + else: + # Reuse previous prediction (DiT step-skipping cache). + assert _cached_flow_pred is not None, "DiT cache: no cached prediction available for step skip." + flow_pred = _cached_flow_pred + flow_pred_action = _cached_flow_pred_action latents = (noisy_input, noisy_input_action) t = (video_timestep, action_timestep) @@ -811,9 +976,12 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): # on the next inference request after websocket reset/session switch. if extra_args.get("reset", False): state.reset() - # Auto-reset based on model state (before accumulation) - if state.should_reset(text_tokens, 0, self.transformer.local_attn_size): - state.reset() + else: + reset_reason = state.reset_reason(text_tokens, 0, self.transformer.local_attn_size) + if reset_reason == "session": + state.reset() + elif reset_reason == "inference": + state.reset_inference_state() state.language = text_tokens # Frame accumulation: stitched single frame → multi-frame video @@ -959,6 +1127,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): video_out = torch.cat([image, video_out], dim=1) state.current_start_frame += self.num_frame_per_block + state.append_video_latents(video_out) + # q99 denorm: [-1,1] → real values action_out = self._denormalize_action(action_out.float(), embodiment_name) diff --git a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py index 9c1e7873d..5815d93d4 100644 --- a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py +++ b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py @@ -23,7 +23,8 @@ class DreamZeroState: Lifecycle: - Created once in DreamZeroPipeline.__init__() - Mutated every forward() call (frame append, KV cache grow) - - reset() on new session / language change / local_attn_size exceeded + - reset() on new session / language change + - reset_inference_state() on local_attn_size exceeded (KV only) """ def __init__(self) -> None: @@ -65,50 +66,106 @@ class DreamZeroState: self.call_count += 1 return np.stack(frames, axis=0) # (T, H, W, C) + def append_video_latents(self, video_out: torch.Tensor) -> None: + """Append one AR chunk of normalized video latents for later decode. + + Args: + video_out: ``(B, T, C, H, W)`` tensor returned by the denoise loop + (before the ``transpose(1, 2)`` stored in ``DiffusionOutput``). + """ + if video_out.dim() != 5: + raise ValueError(f"Expected 5D video_out, got shape {tuple(video_out.shape)}") + # Upstream ``torch.cat(self.video_across_time, dim=2)`` uses (B, C, T, H, W). + chunk = video_out.transpose(1, 2).detach().cpu() + self.video_latents_across_time.append(chunk) + logger.info( + "append_video_latents: chunk_shape=%s total_chunks=%d total_latent_t=%d", + tuple(chunk.shape), + len(self.video_latents_across_time), + int(sum(c.shape[2] for c in self.video_latents_across_time)), + ) + + def get_concatenated_video_latents(self) -> torch.Tensor | None: + """Return all accumulated chunks concatenated along the time dimension.""" + if not self.video_latents_across_time: + return None + if len(self.video_latents_across_time) == 1: + return self.video_latents_across_time[0] + return torch.cat(self.video_latents_across_time, dim=2) + + def clear_video_latents(self) -> None: + """Drop accumulated video latents without resetting KV/frame state.""" + self.video_latents_across_time = [] + # ------------------------------------------------------------------ # Reset / should_reset # ------------------------------------------------------------------ - def reset(self) -> None: - """Clear all state.""" - self.stitched_buffer: deque[np.ndarray] = deque(maxlen=FRAMES_PER_CHUNK) - self.call_count: int = 0 + def reset(self, *, clear_video_latents: bool = True) -> None: + """Clear session state. + + Args: + clear_video_latents: When ``False``, keep ``video_latents_across_time`` + so offline/online export can decode the full AR rollout. + """ + saved_video_latents = [] if clear_video_latents else list(self.video_latents_across_time) + self.stitched_buffer = deque(maxlen=FRAMES_PER_CHUNK) + self.call_count = 0 + # Normalized VAE latents per AR chunk, each (B, C, T, H, W) on CPU. + # Concatenated along T before decode (matches upstream video_across_time). + self.video_latents_across_time = saved_video_latents # KV cache once robot-policy diffusion supports that integration. - self.kv_cache: list[torch.Tensor] | None = None - self.kv_cache_neg: list[torch.Tensor] | None = None - self.crossattn_cache: list[dict[str, bool | torch.Tensor | None]] | None = None - self.crossattn_cache_neg: list[dict[str, bool | torch.Tensor | None]] | None = None - self.current_start_frame: int = 0 + self.kv_cache = None + self.kv_cache_neg = None + self.crossattn_cache = None + self.crossattn_cache_neg = None + self.current_start_frame = 0 - self.clip_feas: torch.Tensor | None = None - self.ys: torch.Tensor | None = None - self.language: torch.Tensor | None = None + self.clip_feas = None + self.ys = None + if clear_video_latents: + self.language = None - def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool: - """Determine if state should be reset before this forward().""" + def reset_inference_state(self) -> None: + """Reset KV/frame state after local attention rolls without dropping video latents.""" + self.reset(clear_video_latents=False) + + def reset_reason( + self, + text_tokens: torch.Tensor | None, + num_video_frames: int, + local_attn_size: int, + ) -> str | None: + """Return why state should reset before the next forward(), if any.""" if self.language is None: logger.info("language is None, resetting") - return True + return "session" if text_tokens is not None and not torch.equal(self.language, text_tokens): logger.info("language changed, resetting") - return True + return "session" # NOTE: after accumulate_frames, num_video_frames is the accumulated T # (1 for first call, 4 for subsequent). Only reset on true single-frame # which happens when the stitched_buffer was cleared externally. if num_video_frames == 1 and self.call_count > 1: logger.info("single frame input after first call, resetting") - return True + return "session" if local_attn_size != -1 and self.current_start_frame >= local_attn_size: logger.info( - "current_start_frame %d >= local_attn_size %d, resetting", self.current_start_frame, local_attn_size + "current_start_frame %d >= local_attn_size %d, resetting inference state", + self.current_start_frame, + local_attn_size, ) - return True + return "inference" - return False + return None + + def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool: + """Determine if state should be reset before this forward().""" + return self.reset_reason(text_tokens, num_video_frames, local_attn_size) is not None # ------------------------------------------------------------------ # KV cache management diff --git a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py index 20c8cd8cd..d039e53c8 100644 --- a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py +++ b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py @@ -9,13 +9,30 @@ import torch class DreamZeroVideoExportWorkerExtension: """DreamZero worker RPCs used by offline example video export.""" + @staticmethod + def _latents_to_uint8_frames(decoded: torch.Tensor) -> torch.Tensor: + decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous() + decoded = decoded.clamp(-1, 1) * 0.5 + 0.5 + return (decoded * 255.0).round().to(torch.uint8).cpu() + def decode_video_latents_to_uint8(self, video_latents: torch.Tensor) -> torch.Tensor: if self.model_runner is None or self.model_runner.pipeline is None: raise RuntimeError("DreamZero pipeline is not initialized on this worker.") with torch.inference_mode(): decoded = self.model_runner.pipeline.decode_video_latents(video_latents) - decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous() - decoded = decoded.clamp(-1, 1) * 0.5 + 0.5 - decoded = (decoded * 255.0).round().to(torch.uint8).cpu() - return decoded + return self._latents_to_uint8_frames(decoded) + + def decode_accumulated_video_latents_to_uint8(self, session_id: str = "default") -> torch.Tensor: + """Decode all AR-chunk latents accumulated on the server for one session.""" + if self.model_runner is None or self.model_runner.pipeline is None: + raise RuntimeError("DreamZero pipeline is not initialized on this worker.") + + with torch.inference_mode(): + decoded = self.model_runner.pipeline.decode_accumulated_video_latents(session_id) + return self._latents_to_uint8_frames(decoded) + + def clear_accumulated_video_latents(self, session_id: str = "default") -> None: + if self.model_runner is None or self.model_runner.pipeline is None: + raise RuntimeError("DreamZero pipeline is not initialized on this worker.") + self.model_runner.pipeline.clear_accumulated_video_latents(session_id) diff --git a/vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py b/vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py new file mode 100644 index 000000000..57d9b866c --- /dev/null +++ b/vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Make diffusers Wan VAE feat_cache compile-friendly. + +diffusers ``WanResample`` stores the string sentinel ``"Rep"`` in feat_cache slots. +``tensor == "Rep"`` returns a non-Tensor and breaks ``torch.compile(fullgraph=True)``. +We replace it with a scalar int8 tensor marker (value 0). +""" + +from __future__ import annotations + +from typing import Any + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_PATCHED = False + + +def _is_feat_rep(entry: Any) -> bool: + if not isinstance(entry, torch.Tensor): + return False + return entry.numel() == 1 and entry.dtype == torch.int8 + + +def _make_feat_rep(x: torch.Tensor) -> torch.Tensor: + return torch.zeros(1, device=x.device, dtype=torch.int8) + + +def _wan_resample_forward_tensorized(self, x, feat_cache=None, feat_idx=[0]): + """``WanResample.forward`` with tensor feat_cache sentinel (no ``"Rep"`` string).""" + from diffusers.models.autoencoders.autoencoder_kl_wan import CACHE_T + + b, c, t, h, w = x.size() + if self.mode == "upsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = _make_feat_rep(x) + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and not _is_feat_rep(feat_cache[idx]): + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and _is_feat_rep(feat_cache[idx]): + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) + if _is_feat_rep(feat_cache[idx]): + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) + x = self.resample(x) + x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4) + + if self.mode == "downsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + +def apply_wan_vae_feat_cache_tensor_patch() -> None: + """Monkeypatch diffusers WanResample once per process.""" + global _PATCHED + if _PATCHED: + return + + from diffusers.models.autoencoders.autoencoder_kl_wan import WanResample + + WanResample.forward = _wan_resample_forward_tensorized + _PATCHED = True + logger.info("DreamZero: applied Wan VAE feat_cache tensor sentinel patch.") diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 278f1df40..288973637 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -168,8 +168,17 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): # Apply torch.compile if not in eager mode if not self.od_config.enforce_eager: if current_omni_platform.supports_torch_inductor(): - self._compile_transformer("transformer") - self._compile_transformer("transformer_2") + if hasattr(self.pipeline, "setup_compile"): + try: + self.pipeline.setup_compile() + except Exception as exc: + logger.warning( + "Model runner: setup_compile() failed (%s); running without compile.", + exc, + ) + else: + self._compile_transformer("transformer") + self._compile_transformer("transformer_2") else: logger.warning( "Model runner: Platform %s does not support torch inductor, skipping torch.compile.", @@ -306,8 +315,15 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): # num_inference_steps at all (i.e., just resets state and clears # stale residuals). num_inference_steps = req.sampling_params.num_inference_steps - if self.od_config.cache_backend == "tea_cache" and num_inference_steps is None: - num_inference_steps = 0 + if num_inference_steps is None and self.od_config.cache_backend in ( + "tea_cache", + "step_cache", + ): + # TeaCache refresh ignores the value; step_cache refresh is a + # no-op (per-chunk state resets in the denoise loop). DreamZero often + # leaves sampling_params.num_inference_steps unset and uses the + # pipeline default instead. + num_inference_steps = getattr(self.pipeline, "num_inference_steps", 0) or 0 if num_inference_steps is not None: self.cache_backend.refresh(self.pipeline, num_inference_steps) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 721e58c7b..ac6140291 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -520,7 +520,7 @@ class AsyncOmniEngine: return mm_data = prompt.get("multi_modal_data") - if not isinstance(mm_data, Mapping) or not mm_data: + if not isinstance(mm_data, dict) or not mm_data: return from vllm.multimodal.hasher import MultiModalHasher @@ -842,6 +842,14 @@ class AsyncOmniEngine: "mag_max_skip_steps": 5, "mag_retention_ratio": 0.1, } + if cache_backend in ("step_cache"): + return { + "step_cache_dit_enabled": True, + "velocity_sim_thresholds": [0.95, 0.93], + "velocity_skip_countdowns": [4, 2], + "step_cache_dit_min_history": 2, + "step_cache_dit_max_history": 2, + } return None @staticmethod diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 905ff4ede..984706ae0 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -512,7 +512,7 @@ class OmniServeCommand(CLISubcommand): "--cache-backend", type=str, default="none", - help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache'", + help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache', 'step_cache'", ) omni_config_group.add_argument( "--cache-config", diff --git a/vllm_omni/entrypoints/openpi/connection.py b/vllm_omni/entrypoints/openpi/connection.py index 69f722ac8..b65264583 100644 --- a/vllm_omni/entrypoints/openpi/connection.py +++ b/vllm_omni/entrypoints/openpi/connection.py @@ -8,7 +8,13 @@ Protocol (compatible with OpenPI policy clients): Infer -> client sends msgpack(obs), server sends msgpack(ndarray) Reset -> client sends msgpack({endpoint:reset}), server sends msgpack(status) -NumPy values use the msgpack-numpy marker mapping: +NumPy values use msgpack-numpy marker mappings. Outbound payloads follow the +openpi-client format so official OpenPI clients can decode action responses: + ndarray -> {__ndarray__: true, dtype, shape, data} + scalar -> {__npgeneric__: true, dtype, data} + +Inbound payloads accept both the openpi-client markers above and the legacy +vLLM-native markers: ndarray -> {nd: true, type, kind, shape, data} scalar -> {nd: false, type, kind, data} """ @@ -41,18 +47,16 @@ def _pack_numpy(obj: Any) -> Any: if not obj.flags.c_contiguous: obj = np.ascontiguousarray(obj) return { - b"nd": True, + b"__ndarray__": True, b"data": obj.tobytes(), - b"type": obj.dtype.str, - b"kind": obj.dtype.kind, + b"dtype": obj.dtype.str, b"shape": obj.shape, } if isinstance(obj, np.generic): return { - b"nd": False, - b"data": obj.tobytes(), - b"type": obj.dtype.str, - b"kind": obj.dtype.kind, + b"__npgeneric__": True, + b"data": obj.item(), + b"dtype": obj.dtype.str, } raise TypeError(f"Unsupported type: {type(obj)!r}") @@ -67,27 +71,67 @@ def _decode_marker_text(value: Any) -> str: return str(value) +def _is_truthy_marker(value: Any) -> bool: + return value is True or value == 1 + + +def _decode_vllm_numpy_marker(obj: dict[Any, Any]) -> Any: + nd = _mapping_get(obj, "nd", _MISSING) + dtype = _mapping_get(obj, "type", _MISSING) + kind = _mapping_get(obj, "kind", _MISSING) + data = _mapping_get(obj, "data", _MISSING) + if nd is _MISSING or dtype is _MISSING or kind is _MISSING or data is _MISSING: + return _MISSING + + dtype_obj = np.dtype(_decode_marker_text(dtype)) + kind_text = _decode_marker_text(kind) + if dtype_obj.kind != kind_text: + raise ValueError(f"NumPy dtype marker kind mismatch: {dtype_obj.kind!r} != {kind_text!r}") + if dtype_obj.kind in ("V", "O", "c"): + raise ValueError(f"Unsupported dtype: {dtype_obj}") + + array = np.frombuffer(data, dtype=dtype_obj).copy() + if nd: + shape = _mapping_get(obj, "shape", _MISSING) + if shape is _MISSING: + raise ValueError("NumPy ndarray marker is missing shape") + return array.reshape(tuple(shape)) + return array[0] + + +def _decode_openpi_numpy_marker(obj: dict[Any, Any]) -> Any: + if _is_truthy_marker(_mapping_get(obj, "__ndarray__", False)): + data = _mapping_get(obj, "data", _MISSING) + dtype = _mapping_get(obj, "dtype", _MISSING) + shape = _mapping_get(obj, "shape", _MISSING) + if data is _MISSING or dtype is _MISSING or shape is _MISSING: + raise ValueError("OpenPI ndarray marker is missing required fields") + + dtype_obj = np.dtype(_decode_marker_text(dtype)) + if dtype_obj.kind in ("V", "O", "c"): + raise ValueError(f"Unsupported dtype: {dtype_obj}") + array = np.frombuffer(data, dtype=dtype_obj).copy() + return array.reshape(tuple(shape)) + + if _is_truthy_marker(_mapping_get(obj, "__npgeneric__", False)): + data = _mapping_get(obj, "data", _MISSING) + dtype = _mapping_get(obj, "dtype", _MISSING) + if data is _MISSING or dtype is _MISSING: + raise ValueError("OpenPI scalar marker is missing required fields") + dtype_obj = np.dtype(_decode_marker_text(dtype)) + if dtype_obj.kind in ("V", "O", "c"): + raise ValueError(f"Unsupported dtype: {dtype_obj}") + return dtype_obj.type(data) + + return _MISSING + + def _unpack_numpy(obj: Any) -> Any: if isinstance(obj, dict): - nd = _mapping_get(obj, "nd", _MISSING) - dtype = _mapping_get(obj, "type", _MISSING) - kind = _mapping_get(obj, "kind", _MISSING) - data = _mapping_get(obj, "data", _MISSING) - if nd is not _MISSING and dtype is not _MISSING and kind is not _MISSING and data is not _MISSING: - dtype_obj = np.dtype(_decode_marker_text(dtype)) - kind_text = _decode_marker_text(kind) - if dtype_obj.kind != kind_text: - raise ValueError(f"NumPy dtype marker kind mismatch: {dtype_obj.kind!r} != {kind_text!r}") - if dtype_obj.kind in ("V", "O", "c"): - raise ValueError(f"Unsupported dtype: {dtype_obj}") - - array = np.frombuffer(data, dtype=dtype_obj).copy() - if nd: - shape = _mapping_get(obj, "shape", _MISSING) - if shape is _MISSING: - raise ValueError("NumPy ndarray marker is missing shape") - return array.reshape(tuple(shape)) - return array[0] + for decoder in (_decode_vllm_numpy_marker, _decode_openpi_numpy_marker): + decoded = decoder(obj) + if decoded is not _MISSING: + return decoded return {key: _unpack_numpy(value) for key, value in obj.items()} if isinstance(obj, list): return [_unpack_numpy(value) for value in obj] From 208932210f118940cd88880c3f7ab52960903327 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Fri, 19 Jun 2026 00:05:26 -0700 Subject: [PATCH 14/40] [Chore][TTS] Reference recipes.vllm.ai recipes, remove dead code, and fix stale references (#4567) Signed-off-by: Yueqian Lin --- README.md | 1 + docs/models/supported_models.md | 5 +- docs/serving/speech_api.md | 7 + recipes/README.md | 11 +- .../entrypoints/openai/serving_speech.py | 1 - .../entrypoints/openai/tts_adapters/base.py | 3 - .../models/cosyvoice3/cosyvoice3.py | 35 +---- .../models/cosyvoice3/cosyvoice3_code2wav.py | 73 ++++------- .../models/cosyvoice3/pipeline.py | 6 +- .../model_executor/models/cosyvoice3/utils.py | 22 ---- .../models/covo_audio/config_covo_audio.py | 1 - .../models/covo_audio/token2wav.py | 5 +- .../fish_speech/fish_speech_dac_decoder.py | 18 ++- .../models/fish_speech/fish_speech_fast_ar.py | 11 +- .../models/fish_speech/prompt_utils.py | 10 -- .../models/glm_tts/glm_tts_dit.py | 4 - .../configuration_higgs_audio_v2.py | 12 -- .../higgs_audio_v2/higgs_audio_v2_talker.py | 35 +++-- .../models/higgs_audio_v3/ref_audio_cache.py | 111 ---------------- .../models/ming_flash_omni/projectors.py | 27 ---- .../models/ming_flash_omni/text_processing.py | 91 ------------- .../models/ming_tts/config_ming_tts.py | 2 - .../models/ming_tts/constants.py | 1 - .../models/ming_tts/ming_tts.py | 4 - .../models/ming_tts/ming_tts_llm.py | 14 +- .../models/moss_tts/audio_tokenizer.py | 2 +- .../moss_tts/modeling_moss_tts_talker.py | 4 +- .../configuration_moss_tts_nano.py | 60 --------- .../moss_tts_nano/modeling_moss_tts_nano.py | 1 - .../models/omnivoice/omnivoice.py | 3 +- .../models/omnivoice/omnivoice_generator.py | 23 ++-- .../qwen3_tts/configuration_qwen3_tts.py | 11 +- .../models/qwen3_tts/prompt_embeds_builder.py | 8 +- .../models/qwen3_tts/qwen3_tts_talker.py | 7 - .../modeling_qwen3_tts_tokenizer_v2.py | 5 - .../qwen3_tts/tokenizer_25hz/vq/core_vq.py | 5 - .../models/step_audio2/step_audio2.py | 8 +- .../step_audio2/step_audio2_constants.py | 36 ------ .../models/step_audio2/step_audio2_thinker.py | 30 +---- .../step_audio2/step_audio2_token2wav.py | 121 ------------------ .../models/voxcpm2/minicpm4_paged.py | 17 +-- .../models/voxcpm2/voxcpm2_talker.py | 43 +++---- .../models/voxtral_tts/voxtral_tts.py | 3 - .../voxtral_tts_audio_generation.py | 26 +--- .../voxtral_tts_audio_tokenizer.py | 4 +- .../transformers_utils/processors/ming.py | 1 - 46 files changed, 141 insertions(+), 787 deletions(-) delete mode 100644 vllm_omni/model_executor/models/higgs_audio_v3/ref_audio_cache.py delete mode 100644 vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py diff --git a/README.md b/README.md index cf1658b45..e300b6b8d 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,7 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn - [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/) - [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/) - [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/) +- [Deployment Recipes](https://recipes.vllm.ai) for vLLM-Omni model serving ## Contributing diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 119f3bf68..e69f462cf 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -6,6 +6,8 @@ vLLM-Omni supports unified multimodal comprehension and generation models across If vLLM-Omni natively supports a model, its implementation can be found in and . +For deployment recipes (TTS and multimodal), see [recipes.vllm.ai](https://recipes.vllm.ai) and the in-repo [`recipes/`](https://github.com/vllm-project/vllm-omni/tree/main/recipes) directory. + ## List of Supported Models