Compare commits

...
Author SHA1 Message Date
Hongsheng Liu dc31c49b76 Revert "[Core] Make diffusion worker titles topology-aware (#6307)"
This reverts commit 13173d6bbb.
2026-08-28 17:29:25 +08:00
Wallbreazzz 4ad455e61e [Feature][TTS] adaptive chunk ramp (Phase 2 buffer-feedback controller) (#6001)
Signed-off-by: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com>
2026-08-28 02:11:05 -07:00
amy-why-3459 da94c3ebe1 [Bugfix][MiniCPM-o] Bind omni-duplex video frames to the unit they close (#6404)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-28 16:58:20 +08:00
prettygirlisnotme 13173d6bbb [Core] Make diffusion worker titles topology-aware (#6307)
Signed-off-by: Dragonborn <1078757180@qq.com>
2026-08-28 16:00:30 +08:00
Sy03 009b992c5b [Bugfix][MiniCPM-o] Isolate duplex handoff failures and barge-in (#6170)
Signed-off-by: Sy03 <1370724210@qq.com>
2026-08-28 15:46:41 +08:00
Mu GuanLinandHongsheng Liu 04b97a5922 [Perf] MiniMax-H3 VAE Decoder Ops (#6607)
Signed-off-by: mglyn <1203789601@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-28 15:14:33 +08:00
TJianandroot 15c6ab8fa7 [AMD][CI] Fix AMD CI (Partial) (#6704)
Signed-off-by: root <root@tw020.mia.tensorwave.lan>
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
Co-authored-by: root <root@tw020.mia.tensorwave.lan>
2026-08-28 14:34:11 +08:00
zzh c201fd77c6 Add HYvideo1.5 benchmark (#6349)
Signed-off-by: zzh <943967662@qq.com>
2026-08-28 14:31:03 +08:00
b6a4baa659 [Feature] Support text encoder Disaggregation for H3 (#5885)
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Signed-off-by: yuanwu <yuan.wu@intel.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: yuanwu <yuan.wu@intel.com>
2026-08-28 14:27:05 +08:00
WeiQing Chen ed837e350e [Test] Stabilize MiniMax-H3 reference accuracy compile (#6688)
Signed-off-by: david6666666 <530634352@qq.com>
2026-08-28 14:18:28 +08:00
BruceLoveDecimalandamy-why-3459 ab8d45d3cc [Bugfix][MiniCPM-o] Retire CFM DiT CUDA graphs a generation at a time (#6587)
Signed-off-by: BruceLoveDecimal <liuqihao970610@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-28 14:07:57 +08:00
Eugene Zhao f8f8aa8eb7 [Bugfix] Honor stage selection in online profiler endpoints (#6609)
Signed-off-by: qujing226 <qujing226@gmail.com>
2026-08-28 11:41:47 +08:00
Hongsheng Liu 6ec5ab11bd [Bugfix] Make HWR filesystem lifecycle transitions fail-closed (#6692)
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-08-28 09:30:17 +08:00
Hongsheng Liuandhsliuustc0106 4245fd6955 [Fix] Restore MammothModa2 support (#6694)
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
2026-08-28 08:32:32 +08:00
Mu GuanLinandHongsheng Liu 5e4daaf197 [Diffusion] Add LTX-2.5 Diffusion VAE decoder support (#6189)
Signed-off-by: mglyn <1203789601@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-27 23:54:40 +08:00
NATUREandamy-why-3459 262ff14693 [Bugfix][MiniCPM-o] Bound native auto-response continuation (#6630)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-27 19:51:47 +08:00
Deep Shahandamy-why-3459 069e6c5fbc [CI/Build][MiniCPM-o] Derive duplex admission-probe limit from the deploy config (#6678)
Signed-off-by: Deep Shah <deep@socratic.co>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-27 19:20:01 +08:00
NATURE 54c1fc64b7 [Bugfix] Gate async-chunk segment resume by session mode (#6680)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
2026-08-27 18:52:09 +08:00
Zhanshenzhuiandamy-why-3459 c6df39e9f7 [Benchmark][MiniCPM-o] Add Omni-DuplexEval support (#6634)
Signed-off-by: Zhanshenzhui <99891002+zyforsure@users.noreply.github.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-27 16:31:50 +08:00
Zheng Wengangandamy-why-3459 94b1546ce3 [Model] Qwen3-Omni: register thinker-only pipeline for Instruct serve (#6284)
Signed-off-by: ZhengWG <zwg0606@gmail.com>
Signed-off-by: Zheng Wengang <zwg0606@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-27 08:21:03 +00:00
19da238246 [Bugfix][NPU] Eager-import mindiesd for every diffusion attention backend (#6054)
Signed-off-by: brandneway <gyuan4892@gmail.com>
Co-authored-by: changetheway <guotaoyuan1@h-partners.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
2026-08-27 14:41:36 +08:00
Zhou TaichangandClaude Opus 5 651b03225b [Rebase] Rebase to vllm 0.28.0 (#6606)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-27 06:12:55 +00:00
Hongsheng Liuandhsliuustc0106 da4a08b65b [Diffusion] Add FLUX.2-klein BF16 HWR contract (#6651)
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
2026-08-27 12:51:15 +08:00
SYLARandHongsheng Liu 0e44fc7406 [diffusion][feature] Add offline SVDQuant W4A4 support (#6162)
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-27 08:45:53 +08:00
Sy03andsy03 432b2c6ca3 [Model] Add native full-duplex Nemotron VoiceChat serving (#6089)
Signed-off-by: Sysy <1370724210@qq.com>
Signed-off-by: sy03 <137072421@qq.com>
Co-authored-by: sy03 <137072421@qq.com>
2026-08-26 13:19:20 -07:00
c2723979c0 [Bugfix] Validate per-stage device layout before spawning workers (#5003) (#5742)
Signed-off-by: ChoHee15 <cc5281@126.com>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-08-26 19:25:22 +00:00
vOv 2d4724eb74 [Bugfix][Diffusion] Fix FLASH_ATTN cross-attention key-padding unpad (#5866)
Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>
2026-08-26 17:20:59 +00:00
Anurag 91dd7fcaec [Bugfix] Reject instructions changes on locked native duplex sessions (#6318)
Signed-off-by: ANURAG KANADE <anuragkanade6@gmail.com>
2026-08-26 16:50:31 +00:00
harley 926e7fbfa2 [Recipe] zai-org/GLM-TTS (2x non-standard RTX 4090 48GB) (#5769)
Signed-off-by: harley <480497125@qq.com>
2026-08-26 22:16:49 +08:00
Yueqian LinandYueqian Lin 8cfede4309 [Perf][Engine] Event-driven orchestration loop (opt-in) — S1 of #4855 (#5221)
Signed-off-by: Yueqian Lin <gao@yueqian.us>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <gao@yueqian.us>
2026-08-26 22:15:25 +08:00
Mu GuanLin d3c990dca8 [Model][Bugfix] Improve LTX audio parity and similarity guards (#6342)
Signed-off-by: mglyn <1203789601@qq.com>
2026-08-26 20:59:10 +08:00
psv666andamy-why-3459 5ec576ab4c [Bugfix] Fix missing Realtime audio completion (#6564)
Signed-off-by: psv666 <2693925048@qq.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-26 18:27:25 +08:00
WeiQing Chen 030fccb2bf [Test] Add MiniMax-H3 DLO DP2 T2VA smoke (#6555)
Signed-off-by: david6666666 <530634352@qq.com>
2026-08-26 17:17:10 +08:00
Zhichao Zhang 816335cb46 [Doc] Add MiMo-Audio recipe for RTX 5090/5090D 32GB (#6559)
Signed-off-by: chaosansui <zzc15560846421@163.com>
2026-08-26 16:26:24 +08:00
Mu GuanLinandHongsheng Liu 59cf3add40 [Feature] Support MiniMax-H3 Turbo LoRA with DLO (#6550)
Signed-off-by: mglyn <1203789601@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-26 16:15:47 +08:00
Gao Han 7439ed41c2 Revert "[Bugfix] Restore name-based model detection for HF cache snapshot paths" (#6642) 2026-08-26 15:02:47 +08:00
Yueqian Lin 7ba5e8a146 [Bugfix] Restore name-based model detection for HF cache snapshot paths (#6624)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-08-25 23:27:35 -07:00
f90e992dc1 [Model] Optimize CosyVoice3 TensorRT stream handoff (#5673)
Signed-off-by: Allen Wu <allenwu2795@gmail.com>
Co-authored-by: TRAE CLI <noreply@bytedance.com>
Co-authored-by: Sy03 <1370724210@qq.com>
Co-authored-by: TRAE CLI <traecli@bytedance.com>
2026-08-25 22:51:46 -07:00
amy-why-3459 5de6f7c37a [Bugfix][MiniCPM-o] Serve native duplex from shipping YAMLs and fence… (#6619)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-26 10:56:51 +08:00
7ae15c74c8 text and image metric (#6150)
Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-08-26 02:54:45 +00:00
wangyuandCursor aeb348d8f8 [CI]Remove DreamID-Omni and MagiHuman support (#6357)
Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: [Your Name] <your.email@example.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-26 01:31:12 +00:00
XiongZhongwei 40dcfdbb43 [Perf][NPU] Enable native GQA & fused AddNorm for MiniMax-H3 encoder (#6040)
Signed-off-by: MarkPoloChina <1228479724@qq.com>
Signed-off-by: XiongZhongwei <1228479724@qq.com>
2026-08-26 09:14:05 +08:00
Hongsheng Liu a708ae5b7a [Bugfix] Reconcile env-var inventory with post-#6217 main drift (#6631)
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-08-26 06:13:28 +08:00
381 changed files with 30075 additions and 8811 deletions

No files matched your search

+89 -20
View File
@@ -1,4 +1,7 @@
#!/bin/bash #!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# vllm-omni customized version # vllm-omni customized version
# Based on: https://github.com/vllm-project/ci-infra/blob/main/buildkite/bootstrap-amd.sh # Based on: https://github.com/vllm-project/ci-infra/blob/main/buildkite/bootstrap-amd.sh
# Last synced: 2025-12-15 # Last synced: 2025-12-15
@@ -6,6 +9,9 @@
set -euo pipefail set -euo pipefail
# The bootstrap runs from the repository root; ShellCheck does not follow
# sourced files unless invoked with -x.
# shellcheck disable=SC1091
source .buildkite/common/scripts/resolve_skip_ci.sh source .buildkite/common/scripts/resolve_skip_ci.sh
if [[ -z "${RUN_ALL:-}" ]]; then if [[ -z "${RUN_ALL:-}" ]]; then
@@ -63,6 +69,8 @@ upload_pipeline() {
# Install minijinja # Install minijinja
ls .buildkite || buildkite-agent annotate --style error 'Please merge upstream main branch for buildkite CI' ls .buildkite || buildkite-agent annotate --style error 'Please merge upstream main branch for buildkite CI'
curl -sSfL https://github.com/mitsuhiko/minijinja/releases/download/2.3.1/minijinja-cli-installer.sh | sh curl -sSfL https://github.com/mitsuhiko/minijinja/releases/download/2.3.1/minijinja-cli-installer.sh | sh
# Installed by the minijinja bootstrap above and only present on the CI agent.
# shellcheck disable=SC1091
source /var/lib/buildkite-agent/.cargo/env source /var/lib/buildkite-agent/.cargo/env
if [[ $BUILDKITE_PIPELINE_SLUG == "fastcheck" ]]; then if [[ $BUILDKITE_PIPELINE_SLUG == "fastcheck" ]]; then
@@ -90,22 +98,82 @@ upload_pipeline() {
cd .buildkite/amd cd .buildkite/amd
# Select test definition file: merge suite for main, ready suite for PRs. # Select test definition file: merge suite for main, ready suite for PRs.
# For debugging, DEBUG_TEST_YAML can override the selection — accepts # For debugging, DEBUG_TEST_YAML accepts a comma-separated list containing
# "merge" or "ready" (case-insensitive). # "merge" and/or "ready" (case-insensitive). Multiple suites are combined
# before rendering so they share one amd-build step.
if [[ -n "${DEBUG_TEST_YAML:-}" ]]; then if [[ -n "${DEBUG_TEST_YAML:-}" ]]; then
case "${DEBUG_TEST_YAML,,}" in declare -a DEBUG_TEST_SPECS=()
merge) declare -A SEEN_DEBUG_TESTS=()
TEST_YAML="test-amd-merge.yml" IFS=',' read -ra REQUESTED_DEBUG_TESTS <<< "${DEBUG_TEST_YAML,,}"
;;
ready) for requested_test in "${REQUESTED_DEBUG_TESTS[@]}"; do
TEST_YAML="test-amd-ready.yml" # Permit readable values such as "ready, merge," and ignore the
;; # empty item produced by a trailing comma.
*) requested_test="${requested_test//[[:space:]]/}"
echo "ERROR: DEBUG_TEST_YAML must be 'merge' or 'ready', got '$DEBUG_TEST_YAML'" >&2 [[ -z "$requested_test" ]] && continue
if [[ -n "${SEEN_DEBUG_TESTS[$requested_test]:-}" ]]; then
echo "ERROR: duplicate DEBUG_TEST_YAML suite '$requested_test'" >&2
exit 1 exit 1
;; fi
esac SEEN_DEBUG_TESTS[$requested_test]=1
echo "DEBUG_TEST_YAML override: using $TEST_YAML"
case "$requested_test" in
ready)
DEBUG_TEST_SPECS+=("READY_TESTS:test-amd-ready.yml")
;;
merge)
DEBUG_TEST_SPECS+=("MERGE_TESTS:test-amd-merge.yml")
;;
*)
echo "ERROR: DEBUG_TEST_YAML entries must be 'merge' or 'ready', got '$requested_test'" >&2
exit 1
;;
esac
done
if [[ ${#DEBUG_TEST_SPECS[@]} -eq 0 ]]; then
echo "ERROR: DEBUG_TEST_YAML did not contain a test suite" >&2
exit 1
elif [[ ${#DEBUG_TEST_SPECS[@]} -eq 1 ]]; then
TEST_YAML="${DEBUG_TEST_SPECS[0]#*:}"
else
TEST_YAML=$(mktemp "${TMPDIR:-/tmp}/amd-debug-tests.XXXXXX.yml")
python - "$TEST_YAML" "${DEBUG_TEST_SPECS[@]}" <<'PY'
import sys
import yaml
output_path, *suite_specs = sys.argv[1:]
combined = {"env": {}, "steps": []}
for suite_spec in suite_specs:
group_name, input_path = suite_spec.split(":", 1)
with open(input_path, encoding="utf-8") as test_file:
suite = yaml.safe_load(test_file)
for name, value in (suite.get("env") or {}).items():
previous = combined["env"].get(name, value)
if previous != value:
raise ValueError(
f"Conflicting environment value for {name}: {previous!r} != {value!r}"
)
combined["env"][name] = value
suite_steps = []
for entry in suite.get("steps") or []:
if "group" in entry:
suite_steps.extend(entry.get("steps") or [])
else:
suite_steps.append(entry)
combined["steps"].append({"group": group_name, "steps": suite_steps})
with open(output_path, "w", encoding="utf-8") as output_file:
yaml.safe_dump(combined, output_file, sort_keys=False)
PY
fi
echo "DEBUG_TEST_YAML override: using ${DEBUG_TEST_SPECS[*]}"
elif [[ $BUILDKITE_BRANCH == "main" ]]; then elif [[ $BUILDKITE_BRANCH == "main" ]]; then
TEST_YAML="test-amd-merge.yml" TEST_YAML="test-amd-merge.yml"
else else
@@ -130,19 +198,20 @@ upload_pipeline() {
> pipeline.yaml > pipeline.yaml
) )
cat pipeline.yaml cat pipeline.yaml
if [[ "$TEST_YAML" == "${TMPDIR:-/tmp}/amd-debug-tests."*.yml ]]; then
rm -f -- "$TEST_YAML"
fi
buildkite-agent artifact upload pipeline.yaml buildkite-agent artifact upload pipeline.yaml
buildkite-agent pipeline upload pipeline.yaml buildkite-agent pipeline upload pipeline.yaml
exit 0 exit 0
} }
get_diff() { get_diff() {
$(git add .) git diff --name-only --diff-filter=ACMDR "$(git merge-base origin/main HEAD)"
echo $(git diff --name-only --diff-filter=ACMDR $(git merge-base origin/main HEAD))
} }
get_diff_main() { get_diff_main() {
$(git add .) git diff --name-only --diff-filter=ACMDR HEAD~1
echo $(git diff --name-only --diff-filter=ACMDR HEAD~1)
} }
file_diff=$(get_diff) file_diff=$(get_diff)
@@ -181,7 +250,7 @@ for file in $file_diff; do
# First check if file matches any pattern # First check if file matches any pattern
matches_pattern=0 matches_pattern=0
for pattern in "${patterns[@]}"; do for pattern in "${patterns[@]}"; do
if [[ $file == $pattern* ]] || [[ $file == $pattern ]]; then if [[ $file == "$pattern"* ]] || [[ $file == "$pattern" ]]; then
matches_pattern=1 matches_pattern=1
break break
fi fi
@@ -191,7 +260,7 @@ for file in $file_diff; do
if [[ $matches_pattern -eq 1 ]]; then if [[ $matches_pattern -eq 1 ]]; then
matches_ignore=0 matches_ignore=0
for ignore in "${ignore_patterns[@]}"; do for ignore in "${ignore_patterns[@]}"; do
if [[ $file == $ignore* ]] || [[ $file == $ignore ]]; then if [[ $file == "$ignore"* ]] || [[ $file == "$ignore" ]]; then
matches_ignore=1 matches_ignore=1
break break
fi fi
+44 -16
View File
@@ -21,7 +21,7 @@ steps:
- label: "Simple · Diffusion Test · Shard %N/%t" - label: "Simple · Diffusion Test · Shard %N/%t"
agent_pool: mi300_1 agent_pool: mi300_1
parallelism: 2 parallelism: 4
depends_on: amd-build depends_on: amd-build
mirror_hardwares: [amdproduction] mirror_hardwares: [amdproduction]
grade: Blocking grade: Blocking
@@ -71,26 +71,30 @@ steps:
commands: commands:
- export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn - export VLLM_WORKER_MULTIPROC_METHOD=spawn
- | - pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
bash -c ' - pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
set +e
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
EXIT1=\$?
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)" - label: "Diffusion Sequence Parallelism Test (Need 4 GPUs) · Shard 1/2"
agent_pool: mi300_4 agent_pool: mi300_4
depends_on: amd-build depends_on: amd-build
mirror_hardwares: [amdproduction] mirror_hardwares: [amdproduction]
timeout_in_minutes: 40
grade: Blocking grade: Blocking
commands: commands:
- export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn - export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py - pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
- timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs) · Shard 2/2"
agent_pool: mi300_4
depends_on: amd-build
mirror_hardwares: [amdproduction]
timeout_in_minutes: 40
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
# merge-only tests # merge-only tests
- label: "Diffusion Tensor Parallelism Test" - label: "Diffusion Tensor Parallelism Test"
@@ -175,17 +179,41 @@ steps:
- label: "Qwen3-TTS CustomVoice E2E Test · Shard %N/%t" - label: "Qwen3-TTS CustomVoice E2E Test · Shard %N/%t"
agent_pool: mi300_1 agent_pool: mi300_1
parallelism: 2 parallelism: 3
depends_on: amd-build depends_on: amd-build
mirror_hardwares: [amdproduction] mirror_hardwares: [amdproduction]
grade: Blocking grade: Blocking
timeout_in_minutes: 80
commands: commands:
- | - |
bash -c ' timeout --signal=TERM --kill-after=2m 80m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB TESTS=(
"tests/e2e/online_serving/test_qwen3_tts_customvoice.py::test_default_cuda_graph_startup[async_chunk]"
"tests/e2e/online_serving/test_qwen3_tts_customvoice.py::test_text_to_audio_002[async_chunk]"
"tests/e2e/offline_inference/test_qwen3_tts_customvoice.py::test_text_to_audio_001[no_cuda_graph]"
)
TEST="$${TESTS[$$BUILDKITE_PARALLEL_JOB]}"
pytest -s -v "$$TEST" -m "advanced_model" --run-level "advanced_model"
'
# TODO: Restore this case to the blocking sharded step once its long-running
# Whisper validation on AMD CI is resolved.
- label: "Qwen3-TTS CustomVoice async_chunk Test (quarantined)"
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: NonBlocking
timeout_in_minutes: 80
commands:
- |
timeout --signal=TERM --kill-after=2m 75m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v "tests/e2e/online_serving/test_qwen3_tts_customvoice.py::test_text_to_audio_001[async_chunk]" -m "advanced_model" --run-level "advanced_model"
' '
- label: "Qwen3-TTS Base E2E Test" - label: "Qwen3-TTS Base E2E Test"
+1 -4
View File
@@ -191,10 +191,7 @@ steps:
mirror_hardwares: [amdproduction] mirror_hardwares: [amdproduction]
grade: Blocking grade: Blocking
commands: commands:
- | - pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model"
timeout 20m bash -c '
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model"
'
# TODO: Bagel test on ROCm is very unstable. @tjtanaa # TODO: Bagel test on ROCm is very unstable. @tjtanaa
# Need to debug before reneable numerical changes across large PRs # Need to debug before reneable numerical changes across large PRs
+13 -1
View File
@@ -166,10 +166,10 @@ steps:
- tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
- tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
- vllm_omni/deploy/minicpmo_4_5.yaml - vllm_omni/deploy/minicpmo_4_5.yaml
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
- vllm_omni/experimental/fullduplex/ - vllm_omni/experimental/fullduplex/
- vllm_omni/model_executor/models/minicpmo_4_5/ - vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/models/step_audio2/ - vllm_omni/model_executor/models/step_audio2/
@@ -217,3 +217,15 @@ steps:
pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model' pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
" "
mirror_hardwares: h100_1 mirror_hardwares: h100_1
- label: "Diffusion · HunyuanVideo-1.5 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_hunyuan_video_15_expansion.py
- vllm_omni/diffusion/models/hunyuan_video/
commands:
- |
timeout 40m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
"
mirror_hardwares: h100_1
+23
View File
@@ -272,6 +272,12 @@ steps:
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda" --run-level "full_model" - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_2 mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · HunyuanVideo-1.5 Function Test"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · Doc Test" - label: ":full_moon: Diffusion X2V · Doc Test"
timeout_in_minutes: 60 timeout_in_minutes: 60
commands: commands:
@@ -322,6 +328,23 @@ steps:
exit $$EXIT2 exit $$EXIT2
mirror_hardwares: h100_2 mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · Perf Test · HunyuanVideo-1.5"
key: nightly-diffusion-x2v-performance-hunyuanvideo15
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json
EXIT1=$$?
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json
EXIT2=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
if [ $$EXIT1 -ne 0 ]; then exit $$EXIT1; fi
exit $$EXIT2
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · Perf Test · MiniMax-H3" - label: ":full_moon: Diffusion X2V · Perf Test · MiniMax-H3"
key: nightly-diffusion-x2v-performance-minimax-h3 key: nightly-diffusion-x2v-performance-minimax-h3
timeout_in_minutes: 180 timeout_in_minutes: 180
+31 -1
View File
@@ -205,10 +205,10 @@ steps:
- tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
- tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
- vllm_omni/deploy/minicpmo_4_5.yaml - vllm_omni/deploy/minicpmo_4_5.yaml
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
- vllm_omni/experimental/fullduplex/ - vllm_omni/experimental/fullduplex/
- vllm_omni/model_executor/models/minicpmo_4_5/ - vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/models/step_audio2/ - vllm_omni/model_executor/models/step_audio2/
@@ -307,6 +307,24 @@ steps:
" "
mirror_hardwares: h100_1 mirror_hardwares: h100_1
- label: "Diffusion · MiniMax-H3 DLO DP2 T2VA Test"
timeout_in_minutes: 90
source_file_dependencies:
- tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py
- tests/helpers/assertions.py
- tests/helpers/runtime.py
- vllm_omni/diffusion/models/minimax_h3/
- vllm_omni/diffusion/offloader/
- vllm_omni/diffusion/distributed/
- vllm_omni/entrypoints/openai/
commands:
- |
timeout 80m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py -m 'core_model and cuda and diffusion' --run-level 'core_model'
"
mirror_hardwares: h100_2
- label: "Diffusion · Qwen Image Test" - label: "Diffusion · Qwen Image Test"
source_file_dependencies: source_file_dependencies:
- tests/e2e/online_serving/test_qwen_image.py - tests/e2e/online_serving/test_qwen_image.py
@@ -320,3 +338,15 @@ steps:
pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' --run-level 'core_model' pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' --run-level 'core_model'
" "
mirror_hardwares: h100_1 mirror_hardwares: h100_1
- label: "Diffusion · HunyuanVideo-1.5 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_hunyuan_video_15_expansion.py
- vllm_omni/diffusion/models/hunyuan_video/
commands:
- |
timeout 40m bash -c "
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m 'core_model and diffusion' --run-level 'core_model'
"
mirror_hardwares: h100_1
+4 -4
View File
@@ -105,7 +105,7 @@ For custom loop structures, follow `vllm_omni/diffusion/models/progress_bar.py`
## Path B: Custom/Private Repo Model ## Path B: Custom/Private Repo Model
For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns. For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns.
### B1. Understand the reference repo ### B1. Understand the reference repo
@@ -126,9 +126,9 @@ This is the key design decision for custom models. Follow these placement rules:
| **Pipeline orchestration** (init, forward, denoise loop) | `vllm_omni/diffusion/models/<name>/pipeline_<name>.py` | Always required | | **Pipeline orchestration** (init, forward, denoise loop) | `vllm_omni/diffusion/models/<name>/pipeline_<name>.py` | Always required |
| **Custom transformer/backbone** (ported and adapted to vllm-omni) | `vllm_omni/diffusion/models/<name>/<name>_transformer.py` or similar | `wan2_2.py`, `fusion.py`, `bagel_transformer.py` | | **Custom transformer/backbone** (ported and adapted to vllm-omni) | `vllm_omni/diffusion/models/<name>/<name>_transformer.py` or similar | `wan2_2.py`, `fusion.py`, `bagel_transformer.py` |
| **Custom sub-models** (VAE, fusion, autoencoder) | `vllm_omni/diffusion/models/<name>/` as separate files | `autoencoder.py`, `fusion.py` | | **Custom sub-models** (VAE, fusion, autoencoder) | `vllm_omni/diffusion/models/<name>/` as separate files | `autoencoder.py`, `fusion.py` |
| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | `dreamid_omni` package via git clone | | **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | original-repo package via git clone |
| **Hardcoded model configs** | Module-level dicts in pipeline file | `VIDEO_CONFIG`, `AUDIO_CONFIG` dicts | | **Hardcoded model configs** | Module-level dicts in pipeline file | `VIDEO_CONFIG`, `AUDIO_CONFIG` dicts |
| **Download/setup script** | `examples/offline_inference/<name>/download_<name>.py` | `download_dreamid_omni.py` | | **Download/setup script** | `examples/offline_inference/<name>/download_<name>.py` | `download_<name>.py` |
| **Custom `model_index.json`** | Generated by download script, placed at model root | Minimal: `{"_class_name": "YourPipeline", ...}` | | **Custom `model_index.json`** | Generated by download script, placed at model root | Minimal: `{"_class_name": "YourPipeline", ...}` |
### B3. Handle external dependencies ### B3. Handle external dependencies
@@ -157,7 +157,7 @@ Copy the essential model files into `vllm_omni/diffusion/models/<name>/` and ada
Custom models have two common patterns for weight loading: Custom models have two common patterns for weight loading:
**Pattern 1: Bypass standard loader** (DreamID-Omni style) **Pattern 1: Bypass standard loader** (eager custom init)
When the original model has complex custom init functions that load weights in `__init__`: When the original model has complex custom init functions that load weights in `__init__`:
@@ -16,20 +16,19 @@ vllm_omni/diffusion/models/wan2_2/
The transformer is loaded separately via `weights_sources` + `load_weights()`. Non-transformer components (VAE, text encoder) are loaded in `__init__` via `from_pretrained()`. The transformer is loaded separately via `weights_sources` + `load_weights()`. Non-transformer components (VAE, text encoder) are loaded in `__init__` via `from_pretrained()`.
### Custom model with external deps (e.g., DreamID-Omni) ### Custom model with external deps
``` ```
vllm_omni/diffusion/models/dreamid_omni/ vllm_omni/diffusion/models/<name>/
├── __init__.py # Exports pipeline only ├── __init__.py # Exports pipeline only
├── pipeline_dreamid_omni.py # Pipeline: loads ALL weights in __init__ via custom helpers ├── pipeline_<name>.py # Pipeline: loads ALL weights in __init__ via custom helpers
├── fusion.py # Custom fusion architecture (video + audio cross-attention) └── <backbone>.py # Custom transformer / fusion architecture
└── wan2_2.py # Re-implemented Wan backbone with split API
examples/offline_inference/x_to_video_audio/ examples/offline_inference/<task>/
└── download_dreamid_omni.py # Downloads weights from 3 HF repos + clones code repo └── download_<name>.py # Downloads weights and/or clones code repo
``` ```
All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency (`dreamid_omni` package) imported with try/except. All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency imported with try/except.
### Custom model with ported code (e.g., BAGEL) ### Custom model with ported code (e.g., BAGEL)
@@ -76,7 +75,7 @@ self.weights_sources = [
] ]
``` ```
### Pattern 3: Fully custom loading (DreamID-Omni) ### Pattern 3: Fully custom loading
``` ```
init → load ALL weights eagerly via custom helpers → load_weights() = no-op init → load ALL weights eagerly via custom helpers → load_weights() = no-op
@@ -108,8 +107,8 @@ Standard diffusers `model_index.json`:
Custom model `model_index.json` (minimal): Custom model `model_index.json` (minimal):
```json ```json
{ {
"_class_name": "DreamIDOmniPipeline", "_class_name": "YourCustomPipeline",
"fusion": "DreamID-Omni/dreamid_omni.safetensors" "fusion": "weights/model.safetensors"
} }
``` ```
@@ -117,7 +116,7 @@ The only **required** field is `_class_name` — it must match a key in `_DIFFUS
## External Dependency Management ## External Dependency Management
### Git clone + .pth injection (DreamID-Omni pattern) ### Git clone + .pth injection
```python ```python
def download_dependency(): def download_dependency():
@@ -163,7 +162,7 @@ The engine checks `isinstance(pipeline, SupportImageInput)` at startup to config
Diffusers models use `config.json` in each subfolder. Custom models often use: Diffusers models use `config.json` in each subfolder. Custom models often use:
**Module-level config dicts** (DreamID-Omni): **Module-level config dicts**:
```python ```python
VIDEO_CONFIG = { VIDEO_CONFIG = {
"patch_size": [1, 2, 2], "model_type": "ti2v", "patch_size": [1, 2, 2], "model_type": "ti2v",
@@ -181,7 +180,7 @@ vae_cfg = bagel_cfg.get("vae_config", {})
## Custom Architecture Patterns ## Custom Architecture Patterns
### Split forward API (DreamID-Omni) ### Split forward API
When a fusion model needs to interleave blocks from two backbones: When a fusion model needs to interleave blocks from two backbones:
@@ -0,0 +1,151 @@
"""Benchmark CosyVoice3-style CUDA stream handoff strategies."""
import argparse
import json
import statistics
import time
import torch
from vllm_omni.platforms import current_omni_platform
def run_case(
*,
device: torch.device,
steps: int,
warmup_steps: int,
producer_cycles: int,
estimator_cycles: int,
use_stream_dependencies: bool,
) -> dict[str, float | int | str]:
caller_stream = torch.cuda.Stream(device=device)
estimator_stream = torch.cuda.Stream(device=device)
input_tensor = torch.zeros(1024, device=device)
output_tensor = torch.empty_like(input_tensor)
submission_ms = []
def step(measure: bool) -> None:
with torch.cuda.stream(caller_stream):
torch.cuda._sleep(producer_cycles)
input_tensor.add_(1)
start = time.perf_counter_ns()
if use_stream_dependencies:
estimator_stream.wait_stream(caller_stream)
else:
caller_stream.synchronize()
with torch.cuda.stream(estimator_stream):
torch.cuda._sleep(estimator_cycles)
output_tensor.copy_(input_tensor)
if use_stream_dependencies:
caller_stream.wait_stream(estimator_stream)
else:
estimator_stream.synchronize()
if measure:
submission_ms.append((time.perf_counter_ns() - start) / 1e6)
for _ in range(warmup_steps):
step(measure=False)
caller_stream.synchronize()
estimator_stream.synchronize()
input_tensor.zero_()
torch.accelerator.reset_peak_memory_stats(device)
initial_allocated = torch.accelerator.memory_allocated(device)
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
wall_start = time.perf_counter()
with torch.cuda.stream(caller_stream):
start_event.record()
for _ in range(steps):
step(measure=True)
with torch.cuda.stream(caller_stream):
end_event.record()
end_event.synchronize()
wall_ms = (time.perf_counter() - wall_start) * 1e3
peak_delta_bytes = torch.accelerator.max_memory_allocated(device) - initial_allocated
expected = torch.full_like(output_tensor, steps)
torch.testing.assert_close(output_tensor, expected, rtol=0, atol=0)
return {
"mode": "wait_stream" if use_stream_dependencies else "synchronize",
"steps": steps,
"submission_median_ms": statistics.median(submission_ms),
"submission_p95_ms": statistics.quantiles(submission_ms, n=20)[18],
"submission_mean_ms": statistics.mean(submission_ms),
"gpu_elapsed_ms": start_event.elapsed_time(end_event),
"wall_elapsed_ms": wall_ms,
"peak_allocation_delta_bytes": peak_delta_bytes,
"final_value": int(output_tensor[0].item()),
}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--device", default="cuda:0")
parser.add_argument("--steps", type=int, default=100)
parser.add_argument("--warmup-steps", type=int, default=10)
parser.add_argument("--repeats", type=int, default=7)
parser.add_argument("--producer-cycles", type=int, default=1_000_000)
parser.add_argument("--estimator-cycles", type=int, default=2_000_000)
return parser.parse_args()
def main() -> None:
args = parse_args()
device = torch.device(args.device)
current_omni_platform.set_device(device)
torch.manual_seed(0)
old_results = []
new_results = []
for repeat in range(args.repeats):
common_args = {
"device": device,
"steps": args.steps,
"warmup_steps": args.warmup_steps,
"producer_cycles": args.producer_cycles,
"estimator_cycles": args.estimator_cycles,
}
if repeat % 2 == 0:
old_results.append(run_case(**common_args, use_stream_dependencies=False))
new_results.append(run_case(**common_args, use_stream_dependencies=True))
else:
new_results.append(run_case(**common_args, use_stream_dependencies=True))
old_results.append(run_case(**common_args, use_stream_dependencies=False))
old_submission_ms = statistics.median(result["submission_median_ms"] for result in old_results)
new_submission_ms = statistics.median(result["submission_median_ms"] for result in new_results)
old_gpu_ms = statistics.median(result["gpu_elapsed_ms"] for result in old_results)
new_gpu_ms = statistics.median(result["gpu_elapsed_ms"] for result in new_results)
print(
json.dumps(
{
"device": torch.cuda.get_device_name(device),
"torch": torch.__version__,
"cuda": torch.version.cuda,
"repeats": args.repeats,
"warmup_steps": args.warmup_steps,
"producer_cycles": args.producer_cycles,
"estimator_cycles": args.estimator_cycles,
"old_submission_median_ms": old_submission_ms,
"new_submission_median_ms": new_submission_ms,
"submission_median_speedup": old_submission_ms / new_submission_ms,
"old_gpu_elapsed_median_ms": old_gpu_ms,
"new_gpu_elapsed_median_ms": new_gpu_ms,
"gpu_elapsed_change_percent": (new_gpu_ms - old_gpu_ms) / old_gpu_ms * 100,
"parity": all(result["final_value"] == args.steps for result in old_results + new_results),
"old_runs": old_results,
"new_runs": new_results,
},
indent=2,
)
)
if __name__ == "__main__":
main()
+8 -5
View File
@@ -1,4 +1,4 @@
# Target vllm version: v0.27.0 (releases/v0.27.0 branch). # Target vllm version: v0.28.0 (tag v0.28.0 @ 2cf0a6915ce5).
# Keep VLLM_BASE_TAG aligned with the target vLLM release on every rebase. # Keep VLLM_BASE_TAG aligned with the target vLLM release on every rebase.
# #
# Release-boundary form: there is no precompiled-wheel pin here, because the # Release-boundary form: there is no precompiled-wheel pin here, because the
@@ -14,11 +14,14 @@
# core init), registering libnvJitLink.so.13 with ldconfig, and re-pinning numpy # core init), registering libnvJitLink.so.13 with ldconfig, and re-pinning numpy
# under numba's ceiling. All of it was compensation for a base image and a vLLM # under numba's ceiling. All of it was compensation for a base image and a vLLM
# that disagreed. At a real release boundary they agree, so none of it is # that disagreed. At a real release boundary they agree, so none of it is
# needed. If you pin an untagged commit again, restore the reinstall block — # needed. The 0.28 rebase repeated the cycle: the v0.28.0rc / final-tag wheels
# agent/tasks/20_update_wheel_url.sh fails the run rather than letting the # rode on the v0.27.0 base image with the full reinstall block until the
# Dockerfile silently ship the wrong vLLM. # official v0.28.0 image shipped, at which point this file dropped back to the
# form you are reading. If you pin an untagged commit again, restore the
# reinstall block — agent/tasks/20_update_wheel_url.sh fails the run rather
# than letting the Dockerfile silently ship the wrong vLLM.
ARG VLLM_BASE_IMAGE=vllm/vllm-openai ARG VLLM_BASE_IMAGE=vllm/vllm-openai
ARG VLLM_BASE_TAG=v0.27.0 ARG VLLM_BASE_TAG=v0.28.0
# We use registry-based layer caching for the dependencies, # We use registry-based layer caching for the dependencies,
# where the tag is determined based on the hash of the contents # where the tag is determined based on the hash of the contents
+1 -1
View File
@@ -1,4 +1,4 @@
ARG BASE_IMAGE=vllm/vllm-openai:v0.26.0 ARG BASE_IMAGE=vllm/vllm-openai:v0.28.0
FROM ${BASE_IMAGE} FROM ${BASE_IMAGE}
ARG COMMON_WORKDIR=/app ARG COMMON_WORKDIR=/app
+2 -2
View File
@@ -1,4 +1,4 @@
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.27.0 ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.28.0
FROM ${BASE_IMAGE} AS base FROM ${BASE_IMAGE} AS base
# Declare a variable to know if we want to use the nightly build or the stable build. # Declare a variable to know if we want to use the nightly build or the stable build.
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
# we should swap over to use stable release ASAP. # we should swap over to use stable release ASAP.
# We should avoid relying on custom commits. # We should avoid relying on custom commits.
ARG USE_NIGHTLY_BUILD=0 ARG USE_NIGHTLY_BUILD=0
ARG VLLM_VERSION_OR_COMMIT_HASH=v0.27.0 ARG VLLM_VERSION_OR_COMMIT_HASH=v0.28.0
ARG ARG_PYTORCH_ROCM_ARCH ARG ARG_PYTORCH_ROCM_ARCH
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
+1 -1
View File
@@ -24,7 +24,6 @@ nav:
- Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md - Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
- X-To-Text: user_guide/examples/offline_inference/x_to_text.md - X-To-Text: user_guide/examples/offline_inference/x_to_text.md
- X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
- Online Serving: - Online Serving:
- Image-To-Image: user_guide/examples/online_serving/image_to_image.md - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
- Image-To-Video: user_guide/examples/online_serving/image_to_video.md - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
@@ -54,6 +53,7 @@ nav:
- BitsAndBytes W4: user_guide/quantization/bitsandbytes.md - BitsAndBytes W4: user_guide/quantization/bitsandbytes.md
- MXFP8 W8A8: user_guide/quantization/mxfp8.md - MXFP8 W8A8: user_guide/quantization/mxfp8.md
- MXFP4 W4A4: user_guide/quantization/mxfp4.md - MXFP4 W4A4: user_guide/quantization/mxfp4.md
- SVDQuant W4A4: user_guide/quantization/svdquant.md
- ModelOpt: user_guide/quantization/modelopt.md - ModelOpt: user_guide/quantization/modelopt.md
- AutoRound: user_guide/quantization/autoround.md - AutoRound: user_guide/quantization/autoround.md
- msModelSlim: user_guide/quantization/msmodelslim.md - msModelSlim: user_guide/quantization/msmodelslim.md
+4
View File
@@ -114,6 +114,10 @@ Configuration classes.
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][] - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][] - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
- [vllm_omni.transformers_utils.configs.cosyvoice3.CosyVoice3Config][] - [vllm_omni.transformers_utils.configs.cosyvoice3.CosyVoice3Config][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
- [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][] - [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][]
- [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][] - [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][]
+4 -2
View File
@@ -82,6 +82,8 @@ depends on the installed kernels and model path.
| --- | --- | --- | --- | --- | | --- | --- | --- | --- | --- |
| `SPEAKER_SAMPLES_DIR` | Filesystem path; default `~/.cache/vllm-omni/speakers` | Speech server; read when speaker storage initializes | Environment-only setting. The directory is created; filesystem errors propagate. | Stable | | `SPEAKER_SAMPLES_DIR` | Filesystem path; default `~/.cache/vllm-omni/speakers` | Speech server; read when speaker storage initializes | Environment-only setting. The directory is created; filesystem errors propagate. | Stable |
| `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable | | `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable |
| `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental |
| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; default `0` (off) | Orchestration loop and the serving-side final-output drain; read once when the `Orchestrator` is constructed | Environment-only setting. Values are stripped and case-normalized; any unrecognized value leaves the legacy poll loop selected. | Experimental |
| `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control | | `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control |
| `VLLM_OMNI_ORCH_MONITOR_PATH` | Filesystem path; default `<current-working-directory>/vllm_omni_orch_monitor_<timestamp>.json` | Orchestrator monitor enabled by `--enable-orch-monitor`; read when the monitor is created | Environment-only path override. Parent directories are created; write errors are logged. | Diagnostic | | `VLLM_OMNI_ORCH_MONITOR_PATH` | Filesystem path; default `<current-working-directory>/vllm_omni_orch_monitor_<timestamp>.json` | Orchestrator monitor enabled by `--enable-orch-monitor`; read when the monitor is created | Environment-only path override. Parent directories are created; write errors are logged. | Diagnostic |
| `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` | Float seconds; default `600` | Synchronous Videos API; read when the API server module imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental | | `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` | Float seconds; default `600` | Synchronous Videos API; read when the API server module imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental |
@@ -177,7 +179,7 @@ collection, examples, bug reports, or logs.
## Model-specific variables ## Model-specific variables
The audit found 57 variables read by a single model or pipeline family. They are The audit found 56 variables read by a single model or pipeline family. They are
not listed as public usage options here because doing so would turn implementation not listed as public usage options here because doing so would turn implementation
escape hatches into an accidental compatibility contract. escape hatches into an accidental compatibility contract.
@@ -190,7 +192,7 @@ Every audited model-specific name has a migration disposition in the
| Request scope | 6 | Move request-varying behavior into a declared request-option schema. | | Request scope | 6 | Move request-varying behavior into a declared request-option schema. |
| External | 0 | Retain only when a supported third-party library owns the contract. | | External | 0 | Retain only when a supported third-party library owns the contract. |
| Internalize | 11 | Keep a debug or diagnostic switch out of public documentation and configuration. | | Internalize | 11 | Keep a debug or diagnostic switch out of public documentation and configuration. |
| Deprecate/remove | 7 | Remove a compatibility escape hatch that has no continuing contract. | | Deprecate/remove | 6 | Remove a compatibility escape hatch that has no continuing contract. |
The disposition is a migration target, not a statement that the existing The disposition is a migration target, not a statement that the existing
environment switch is stable. Promote or request-scope work should land in environment switch is stable. Promote or request-scope work should land in
+16 -12
View File
@@ -14,7 +14,7 @@ describe fixed model topology and are not accepted in deploy YAMLs.
Common `PipelineConfig` fields include: Common `PipelineConfig` fields include:
| Field | Description | | Field | Description |
|-------|-------------| | ------- | ------------- |
| `model_type` | Pipeline identifier used during model and config resolution. | | `model_type` | Pipeline identifier used during model and config resolution. |
| `default_deploy_config_name` | Bundled deploy YAML loaded when the user does not pass `deploy_config`. | | `default_deploy_config_name` | Bundled deploy YAML loaded when the user does not pass `deploy_config`. |
| `model_arch` | Default Hugging Face architecture for the pipeline. | | `model_arch` | Default Hugging Face architecture for the pipeline. |
@@ -26,7 +26,7 @@ Common `PipelineConfig` fields include:
Common `StagePipelineConfig` fields include: Common `StagePipelineConfig` fields include:
| Field | Description | | Field | Description |
|-------|-------------| | ------- | ------------- |
| `stage_id` | Stable stage identifier. | | `stage_id` | Stable stage identifier. |
| `model_stage` | Logical stage name used by runtime and strategy resolution. | | `model_stage` | Logical stage name used by runtime and strategy resolution. |
| `execution_type` | `LLM_AR`, `LLM_GENERATION`, or `DIFFUSION`. | | `execution_type` | `LLM_AR`, `LLM_GENERATION`, or `DIFFUSION`. |
@@ -49,17 +49,17 @@ deployment overrides.
The new deploy schema lives under `vllm_omni/deploy/` and is paired with a frozen `PipelineConfig` registered by the model's `pipeline.py`. Each deploy YAML has these top-level fields: The new deploy schema lives under `vllm_omni/deploy/` and is paired with a frozen `PipelineConfig` registered by the model's `pipeline.py`. Each deploy YAML has these top-level fields:
| Field | Type | Required | Default | Description | | Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------| | ------- | ------ | ---------- | --------- | ------------- |
| `base_config` | str (path) | optional | — | Overlay parent (relative or absolute). `stages:` / `platforms:` deep-merged by stage_id; other scalars overlay-wins. Intended for user-authored overlays; prod yamls stay flat. | | `base_config` | str (path) | optional | — | Overlay parent (relative or absolute). `stages:` / `platforms:` deep-merged by stage_id; other scalars overlay-wins. Intended for user-authored overlays; prod yamls stay flat. |
| `async_chunk` | bool | optional | `true` | Enable chunked streaming between stages. Pin to `false` if the pipeline runs end-to-end. | | `async_chunk` | bool | optional | `true` | Enable chunked streaming between stages. Pin to `false` if the pipeline runs end-to-end. |
| `session_mode` | str | optional | `"turn"` | Session behavior. Use `"duplex"` only with a pipeline that enables duplex control. | | `session_mode` | str | optional | `"turn"` | Session behavior. MiniCPM-o 4.5 deploy YAMLs set `"duplex"` so `/v1/realtime` and chat share the same profile. |
| `active_stream_window` | int | optional | `0` | Number of active downstream stream slots; `0` preserves all-stream cycling. | | `active_stream_window` | int | optional | `0` | Number of active downstream stream slots; `0` preserves all-stream cycling. |
| `duplex_session` | dict | optional | runtime defaults | Full-duplex session lifecycle, buffering, replay, and capacity limits. | | `duplex_session` | dict | optional | runtime defaults | Full-duplex session lifecycle, buffering, replay, and capacity limits. |
| `connectors` | dict | optional | `null` | Named connector specs (`{name, extra}`). Referenced by each stage's `input_connectors` / `output_connectors`. See [Connector schema](#connector-schema). | | `connectors` | dict | optional | `null` | Named connector specs (`{name, extra}`). Referenced by each stage's `input_connectors` / `output_connectors`. See [Connector schema](#connector-schema). |
| `edges` | list | optional | `null` | Explicit edge list for the KV transfer graph. Auto-derived from stage inputs if omitted. | | `edges` | list | optional | `null` | Explicit edge list for the KV transfer graph. Auto-derived from stage inputs if omitted. |
| `stages` | list | optional | `[]` | Per-stage runtime overrides matched by `stage_id`. Pipeline stages are still created from `PipelineConfig` when this list is empty. | | `stages` | list | optional | `[]` | Per-stage runtime overrides matched by `stage_id`. Pipeline stages are still created from `PipelineConfig` when this list is empty. |
| `platforms` | dict | optional | `null` | Keyed by `npu` / `rocm` / `xpu`, each contains a `stages:` list with per-platform overrides applied on top of the CUDA defaults. | | `platforms` | dict | optional | `null` | Keyed by `npu` / `rocm` / `xpu`, each contains a `stages:` list with per-platform overrides applied on top of the CUDA defaults. |
| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only`). | | `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only` / `qwen3_omni_moe_thinker_only`). |
| `trust_remote_code` | bool \| null | optional | `null` | **Pipeline-wide.** Trust HF remote code on model load; applies to every stage when specified. | | `trust_remote_code` | bool \| null | optional | `null` | **Pipeline-wide.** Trust HF remote code on model load; applies to every stage when specified. |
| `distributed_executor_backend` | str \| null | optional | `null` | **Pipeline-wide.** Distributed executor backend forwarded to vLLM (`"mp"`, `"ray"`, `"external_launcher"`). If omitted, vLLM auto-selects backend from runtime topology. | | `distributed_executor_backend` | str \| null | optional | `null` | **Pipeline-wide.** Distributed executor backend forwarded to vLLM (`"mp"`, `"ray"`, `"external_launcher"`). If omitted, vLLM auto-selects backend from runtime topology. |
| `dtype` | str \| null | optional | `null` | **Pipeline-wide.** Model dtype for every stage. | | `dtype` | str \| null | optional | `null` | **Pipeline-wide.** Model dtype for every stage. |
@@ -84,7 +84,7 @@ subprocess on one GPU. `ray` / `external_launcher` are not fully supported yet.
Each entry under `stages:` accepts any `StageDeployConfig` field directly (no nested `engine_args:`). Only fields whose value legitimately varies across stages live here; pipeline-wide settings (trust_remote_code, distributed_executor_backend, dtype, quantization, prefix/chunked prefill, DP/PP sizes) are declared at the top level and applied to every stage. Unknown keys fall through to `engine_extras:` and are forwarded to the engine. Frequently used fields are listed below; the source-of-truth schema is `StageDeployConfig` in `vllm_omni/config/stage_config.py`. Each entry under `stages:` accepts any `StageDeployConfig` field directly (no nested `engine_args:`). Only fields whose value legitimately varies across stages live here; pipeline-wide settings (trust_remote_code, distributed_executor_backend, dtype, quantization, prefix/chunked prefill, DP/PP sizes) are declared at the top level and applied to every stage. Unknown keys fall through to `engine_extras:` and are forwarded to the engine. Frequently used fields are listed below; the source-of-truth schema is `StageDeployConfig` in `vllm_omni/config/stage_config.py`.
| Field | Type | Required | Default | Description | | Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------| | ------- | ------ | ---------- | --------- | ------------- |
| `stage_id` | int | required | — | Stage identity; matched against `PipelineConfig.stages[*].stage_id`. | | `stage_id` | int | required | — | Stage identity; matched against `PipelineConfig.stages[*].stage_id`. |
| `max_num_seqs` | int \| null | optional | `null` | Max concurrent sequences per stage. | | `max_num_seqs` | int \| null | optional | `null` | Max concurrent sequences per stage. |
| `gpu_memory_utilization` | float \| null | optional | `null` | Per-stage total memory target; used for automatic KV-cache sizing. | | `gpu_memory_utilization` | float \| null | optional | `null` | Per-stage total memory target; used for automatic KV-cache sizing. |
@@ -94,12 +94,14 @@ Each entry under `stages:` accepts any `StageDeployConfig` field directly (no ne
| `max_num_batched_tokens` | int \| null | optional | `null` | Per-stage prefill/token budget; also contributes to the native maximum in-flight token limit. | | `max_num_batched_tokens` | int \| null | optional | `null` | Per-stage prefill/token budget; also contributes to the native maximum in-flight token limit. |
| `max_model_len` | int \| null | optional | `null` | Per-sequence context or KV length; `-1` enables native cache-capacity auto-fitting, while values above the HF default auto-set `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1`. | | `max_model_len` | int \| null | optional | `null` | Per-sequence context or KV length; `-1` enables native cache-capacity auto-fitting, while values above the HF default auto-set `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1`. |
| `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. | | `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. |
| `devices` | str \| null | optional | `null` | Device list assigned to this stage. | | `devices` | str \| null | optional | `null` | Device list assigned to this stage. The number of device ids must equal this stage's local world size (`tensor_parallel_size` × local data-parallel size × `pipeline_parallel_size`, or `num_replicas` × that product for a replica pool); a mismatch fails early — see the note below. |
| `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_<n>`; values are names registered under top-level `connectors:`. | | `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_<n>`; values are names registered under top-level `connectors:`. |
| `input_connectors` | dict \| null | optional | `null` | Keyed by `from_stage_<n>`; values are names registered under top-level `connectors:`. | | `input_connectors` | dict \| null | optional | `null` | Keyed by `from_stage_<n>`; values are names registered under top-level `connectors:`. |
| `default_sampling_params` | dict \| null | optional | `null` | Baseline sampling params. Deep-merged with pipeline `sampling_constraints` (pipeline wins). | | `default_sampling_params` | dict \| null | optional | `null` | Baseline sampling params. Deep-merged with pipeline `sampling_constraints` (pipeline wins). |
| `engine_extras` | dict | optional | `{}` | Catch-all for engine fields not listed above; deep-merged across overlays and forwarded to the stage engine. | | `engine_extras` | dict | optional | `{}` | Catch-all for engine fields not listed above; deep-merged across overlays and forwarded to the stage engine. |
**Note:** a stage's `devices` count must equal its local world size (`tensor_parallel_size` × `data_parallel_size_local` × `pipeline_parallel_size`, falling back to global `data_parallel_size` when the local size is unset), or `num_replicas` × that product for a replica pool. A mismatch fails early and names the offending stage. A top-level `--tensor-parallel-size` is broadcast to every stage, so it can make a single-GPU stage violate this contract ([issue #5003](gh-issue:5003)); fix that case with `--stage-overrides` (set `tensor_parallel_size` and `devices` together per stage) or set TP only on the multi-GPU stage.
### Connector schema ### Connector schema
Each entry under top-level `connectors:` follows this shape: Each entry under top-level `connectors:` follows this shape:
@@ -113,9 +115,9 @@ connectors:
# Additional connector-specific options # Additional connector-specific options
``` ```
| Connector class | Use case | `extra` keys | | Connector class | Use case | `extra` keys |
|-----------------|----------|--------------| |--------------------------|-----------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------|
| `SharedMemoryConnector` | Same-host KV transfer between stages (default for bundled YAMLs). | None. All payloads use shared memory. | | `SharedMemoryConnector` | Same-host KV transfer between stages (default for bundled YAMLs). | None. All payloads use shared memory. |
| `MooncakeStoreConnector` | Cross-host KV transfer over TCP. Required for multi-node deployments. | `host`, `metadata_server`, `master`, `segment` (int bytes), `localbuf` (int bytes), `proto` (`"tcp"` / `"rdma"`). | | `MooncakeStoreConnector` | Cross-host KV transfer over TCP. Required for multi-node deployments. | `host`, `metadata_server`, `master`, `segment` (int bytes), `localbuf` (int bytes), `proto` (`"tcp"` / `"rdma"`). |
A stage references a connector by name in its `input_connectors` / `output_connectors`: A stage references a connector by name in its `input_connectors` / `output_connectors`:
@@ -135,7 +137,7 @@ stages:
### CLI flags ### CLI flags
| Flag | Description | | Flag | Description |
|------|-------------| | ------ | ------------- |
| `--deploy-config PATH` | Load a deploy YAML. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. | | `--deploy-config PATH` | Load a deploy YAML. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
| `--stage-overrides JSON` | Per-stage JSON overrides, e.g. `'{"0":{"gpu_memory_utilization":0.5}}'`. Per-stage values always win over global flags. | | `--stage-overrides JSON` | Per-stage JSON overrides, e.g. `'{"0":{"gpu_memory_utilization":0.5}}'`. Per-stage values always win over global flags. |
| `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. | | `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. |
@@ -249,7 +251,7 @@ CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
Effective config per stage after the merge: Effective config per stage after the merge:
| Stage | Field | Final value | Source | | Stage | Field | Final value | Source |
|-------|-------|-------------|--------| | ------- | ------- | ------------- | -------- |
| 0 | `gpu_memory_utilization` | `0.9` | base YAML (overlay didn't touch stage 0) | | 0 | `gpu_memory_utilization` | `0.9` | base YAML (overlay didn't touch stage 0) |
| 0 | `max_num_seqs` | `8` | per-stage CLI (`--stage-overrides`) — wins over base `32` | | 0 | `max_num_seqs` | `8` | per-stage CLI (`--stage-overrides`) — wins over base `32` |
| 0 | `max_model_len` | `16384` | global CLI | | 0 | `max_model_len` | `16384` | global CLI |
@@ -273,12 +275,14 @@ when loading a deploy configuration.
Examples: Examples:
For offline inference (assuming the necessary dependencies have been imported): For offline inference (assuming the necessary dependencies have been imported):
```python ```python
model_name = "Qwen/Qwen2.5-Omni-7B" model_name = "Qwen/Qwen2.5-Omni-7B"
omni = Omni(model=model_name, deploy_config="/path/to/deploy_config.yaml") omni = Omni(model=model_name, deploy_config="/path/to/deploy_config.yaml")
``` ```
For online serving: For online serving:
```bash ```bash
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
``` ```
+9
View File
@@ -322,6 +322,15 @@ Each 1-second window records:
On shutdown the server also logs a short summary (`loop_active_pct`, per-replica queue averages/maxima). On shutdown the server also logs a short summary (`loop_active_pct`, per-replica queue averages/maxima).
`loop_idle` and `loop_active` count loop iterations, so their absolute scale is
tied to which orchestration loop is running. Under the default poll loop an idle
orchestrator records roughly one iteration per millisecond. Under the
event-driven loop (`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, see
[Speech API](../serving/speech_api.md#orchestration-loop-experimental)) an idle
orchestrator wakes only once per 0.5 s reconcile timeout, while a busy one still
records one iteration per routed output. Window counts and the `loop_active_pct`
summary are therefore not comparable across the two modes.
### Relationship to other diagnostics ### Relationship to other diagnostics
This monitor is intentionally separate from the existing profiling tools: This monitor is intentionally separate from the existing profiling tools:
+1 -1
View File
@@ -63,7 +63,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
### N-Branch CFG (3+ branches) ### N-Branch CFG (3+ branches)
Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches, DreamID Omni uses 4 branches. Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches.
`predict_noise_with_multi_branch_cfg()` handles these by automatically dispatching N branches across M GPUs using round-robin (rule: branch `i` → rank `i % M`): `predict_noise_with_multi_branch_cfg()` handles these by automatically dispatching N branches across M GPUs using round-robin (rule: branch `i` → rank `i % M`):
+11 -11
View File
@@ -14,12 +14,12 @@ specified in the [Host Weight Runtime module design](../module/host_weight_runti
## Status ## Status
The first implementation provides contracts and a CPU local-filesystem store. The implementation provides contracts and a CPU local-filesystem store. The
The initial diffusion consumer contract additionally defines typed, diffusion consumer additionally defines typed,
representation-independent final-layout identity/restoration mechanics plus a representation-independent final-layout identity/restoration mechanics plus a
concrete BF16-with-preserved-FP32 policy for MiniMax H3. It remains concrete BF16-with-preserved-FP32 policy for MiniMax H3 and
library-only: no loader or DLO path selects, publishes, restores, or transports `black-forest-labs/FLUX.2-klein-4B`. The opt-in no-AllGather DLO integration
that artifact yet. selects, publishes, restores, and transfers these artifacts.
V1 includes: V1 includes:
@@ -33,13 +33,12 @@ V1 includes:
V1 does not include: V1 does not include:
- a public CLI, loader activation, or default-on model integration; - default-on activation or consumers outside no-AllGather DLO;
- lease handoff to DLO or another transport consumer; - online FP8, quantized, merged-adaptation, or additional model producers;
- FP8, quantized, merged-adaptation, or additional model producers; - HWR interaction with DLO AllGather;
- CUDA registration, pinned staging, H2D scheduling, or GPU kernels;
- a remote artifact provider or cross-node coordination; - a remote artifact provider or cross-node coordination;
- automatic eviction; or - automatic eviction; or
- a change to DLO AllGather or no-AllGather behavior. - a change to DLO collective or execution behavior.
## Motivation and use cases ## Motivation and use cases
@@ -255,7 +254,8 @@ The contract is intentionally separate from loader activation:
- `FinalLayoutBF16Producer` accepts only the matching identity context and a - `FinalLayoutBF16Producer` accepts only the matching identity context and a
finalized CPU model. It is `POST_LOAD_ONLY` and `SINGLE_PROCESS` per exact TP finalized CPU model. It is `POST_LOAD_ONLY` and `SINGLE_PROCESS` per exact TP
coordinate. Its BF16 policy preserves model-declared FP32 parameters and coordinate. Its BF16 policy preserves model-declared FP32 parameters and
buffers and revalidates MiniMax H3 mixed-precision invariants. buffers, revalidates MiniMax H3 mixed-precision invariants, and revalidates
FLUX.2-klein's two block stacks, packed QKV mapping, and BF16 base layout.
Other representations reuse source identity, typed parallel identity, tensor Other representations reuse source identity, typed parallel identity, tensor
ownership, and exact restoration only when their policy proves those semantics. ownership, and exact restoration only when their policy proves those semantics.
@@ -220,9 +220,15 @@ model revision and parallel layout with `required`. TP coordinates have distinct
identities, while equivalent DP replicas share them. identities, while equivalent DP replicas share them.
The current producer/consumer boundary is the model-declared final-layout BF16 The current producer/consumer boundary is the model-declared final-layout BF16
contract (MiniMax H3 today). It supports ordinary-loader final layouts for TP1 contract (MiniMax H3 and `black-forest-labs/FLUX.2-klein-4B`). The FLUX.2-klein
and TP2 rank identities plus SP layout identities; online quantization, HSDP, contract covers both transformer block stacks, constructor-stable packed QKV
LoRA/adapted weights, and non-default load formats remain ineligible. HWR mode mapping state, BF16 parameters, and any persistent loader-owned `beta`/`eps`
buffers. The shared HWR machinery supports ordinary-loader final layouts for
TP1 and TP2 rank identities plus SP layout identities. FLUX.2-klein evidence
in this change is TP1-only; its TP2/SP layouts rely on the existing
per-coordinate identity mechanics and remain unmeasured. Online quantization,
HSDP, LoRA/adapted weights, and non-default load formats remain ineligible.
HWR mode
`disabled`, DLO-disabled, and DLO AllGather configurations stop before source `disabled`, DLO-disabled, and DLO AllGather configurations stop before source
identity or store construction and retain the existing checkpoint-mmap or identity or store construction and retain the existing checkpoint-mmap or
ordinary-loader path. ordinary-loader path.
+76 -30
View File
@@ -6,7 +6,7 @@ This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage
- Expose pipeline-level request and latency metrics that span the full multi-stage execution (orchestrator scope). - Expose pipeline-level request and latency metrics that span the full multi-stage execution (orchestrator scope).
- Preserve all upstream vLLM per-engine metrics (`vllm:*`) for stages backed by an AR LLM engine, and reshape their `engine` label into `stage` + `replica` so multi-replica deployments gain per-replica visibility automatically. - Preserve all upstream vLLM per-engine metrics (`vllm:*`) for stages backed by an AR LLM engine, and reshape their `engine` label into `stage` + `replica` so multi-replica deployments gain per-replica visibility automatically.
- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss. - Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss, plus image and diffusion workload, queueing, execution, and memory signals.
- Expose per-replica-edge cross-stage transfer metrics so the slack between E2E latency and the sum of per-stage `gen_time` (queueing, serialization, network) becomes attributable. - Expose per-replica-edge cross-stage transfer metrics so the slack between E2E latency and the sum of per-stage `gen_time` (queueing, serialization, network) becomes attributable.
- Keep the metrics collection overhead low enough that it does not regress TTFA or throughput. - Keep the metrics collection overhead low enough that it does not regress TTFA or throughput.
@@ -20,7 +20,7 @@ vLLM's `unregister_vllm_metrics()` function strips every `prometheus_client` col
### The Problem ### The Problem
vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm:omni_` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm:omni_` / `vllm_omni`. vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm_omni:` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm_omni:` / `vllm_omni`.
Upstream per-engine metrics retain the `vllm:` prefix but are now registered by `OmniPrometheusStatLogger`, a thin subclass of upstream's `PrometheusStatLogger` that reshapes the single `engine` label into a `stage` + `replica` pair (see "OmniPrometheusStatLogger wrap" below). Upstream per-engine metrics retain the `vllm:` prefix but are now registered by `OmniPrometheusStatLogger`, a thin subclass of upstream's `PrometheusStatLogger` that reshapes the single `engine` label into a `stage` + `replica` pair (see "OmniPrometheusStatLogger wrap" below).
@@ -38,7 +38,7 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by
| |
+--------+--------+--------+--------+--------+ +--------+--------+--------+--------+--------+
| | | |
vllm:omni_* vllm:* vllm_omni:* vllm:*
collectors collectors collectors collectors
| | | |
+----+--------+ +-----------+ +----------+ +-----------+ +----+--------+ +-----------+ +----------+ +-----------+
@@ -58,17 +58,30 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by
### Data Flow ### Data Flow
There are four independent paths for metric collection. There are five independent paths for metric collection.
**Path 1: Pipeline-level metrics (`vllm:omni_*`)** **Path 1: Pipeline-level metrics (`vllm_omni:*`)**
`OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress: `OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress:
- `set_running(n)` / `set_waiting(n)` — updated after each request completes. The running count comes from `OmniRequestCounter`, a simple counter incremented/decremented by the Orchestrator as it tracks requests. Waiting is derived as `total - running`. - `set_running(n)` / `set_waiting(n)` — updated after each request completes. The running count comes from `OmniRequestCounter`, a simple counter incremented/decremented by the Orchestrator as it tracks requests. Waiting is derived as `total - running`.
- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm:omni_requests_success_total{finished_reason}`. - `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm_omni:requests_success_total{finished_reason}`.
- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` so a single Counter family covers both natural and aborted completion. - `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` on the existing completion Counter.
- `inc_requests_failed(reason)` — increments the dedicated failure-attribution Counter once per request. Reasons are normalized to the bounded set `client_abort`, `client_disconnect`, `stage_error`, and `unknown`.
**Path 2: Audio modality metrics (`vllm:omni_audio_*`)** **Path 2: Image and diffusion metrics**
Image and diffusion metrics use two cooperating collectors. `OmniPrometheusMetrics` owns pipeline/stage workload families, while `OmniModalityMetrics` owns per-replica diffusion breakdowns. Finished stage messages are consumed once in `OmniBase._process_single_result`, which prevents replayed terminal messages from incrementing Counters or Histograms twice.
- `stage_gen_time_s`, `stage_in_queue_s`, image workload, inference-step, and per-stage peak-memory values are emitted from `StageRequestStats` when a stage finishes. `image_pixels` and `image_count_total` are restricted to image outputs; `num_inference_steps` is restricted to diffusion stages.
- Diffusion engine timings are accumulated per request by `OrchestratorAggregator.accumulate_diffusion_metrics`. Engine-side millisecond values are converted to seconds before `observe_modality_at_finalize(...)` dispatches the per-replica Histogram observations.
- `request_queue_wait_s` is emitted once at final request completion from `pipeline_timings["queue_wait_ms"]`. A present value of `0` is a valid observation; a missing key produces no sample.
- `stage_waiting_requests` is the sum of the latest scheduler snapshots for all live replicas in a stage. AR snapshots arrive through `SchedulerStats`; diffusion snapshots ride normal diffusion output metrics. A dead replica's snapshot is removed. Because this deliberately reuses existing message paths, the gauge can retain its last value while a replica produces no output.
- `kv_wait_s` measures the AR scheduler interval from entering the KV-transfer-free wait state until extraction acknowledgement. The sample is carried to the Orchestrator through the existing engine output and labeled by connector type. Abort/error/removal paths discard incomplete timers rather than emitting partial waits.
Zero and missing values have different meanings. Valid zero-duration queue observations are recorded, while unavailable optional measurements are omitted instead of being exported as synthetic zeroes. `peak_memory_mb`, `diffusion_forward_s`, `diffusion_kv_load_s`, `vae_decode_s`, and `kv_wait_s` therefore appear only when their data source is active.
**Path 3: Audio modality metrics (`vllm_omni:audio_*`)**
`OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites: `OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites:
@@ -76,7 +89,7 @@ There are four independent paths for metric collection.
- `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry. - `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry.
- `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`. - `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`.
**Path 3: Cross-stage transfer metrics (`vllm:omni_transfer_*`)** **Path 4: Cross-stage transfer metrics (`vllm_omni:transfer_*`)**
`OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution. `OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution.
@@ -87,7 +100,7 @@ The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transf
Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected). Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected).
**Path 4: Per-engine metrics (`vllm:*`, stage/replica wrap)** **Path 5: Per-engine metrics (`vllm:*`, stage/replica wrap)**
The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics. The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics.
@@ -97,9 +110,9 @@ The Orchestrator runs in a background thread. The API server (OmniBase) runs in
### Metric Registration and Lifecycle ### Metric Registration and Lifecycle
All `vllm:omni_*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data. All `vllm_omni:*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm:omni_*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime. The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm_omni:*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
## OmniPrometheusStatLogger Wrap ## OmniPrometheusStatLogger Wrap
@@ -130,11 +143,11 @@ A reverse map `(stage_id, replica_id) -> flat_idx` is maintained on the Orchestr
## Gating: `--log-stats` ## Gating: `--log-stats`
All metrics — both the 15 `vllm:omni_*` families and the ~65 upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`. All metrics — both the omni-specific `vllm_omni:*` families and the upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
Behavior with `--log-stats=off` (default): Behavior with `--log-stats=off` (default):
- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm:omni_*` stay materialized but never have data written to them. - The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm_omni:*` stay materialized but never have data written to them.
- `OmniPrometheusStatLogger` is not constructed in `_init_metrics_state`, so the ~65 upstream `vllm:*` wrap families are not registered in the default registry at all. - `OmniPrometheusStatLogger` is not constructed in `_init_metrics_state`, so the ~65 upstream `vllm:*` wrap families are not registered in the default registry at all.
- The engine core's `Scheduler.make_stats()` also short-circuits inside upstream (`if not self.log_stats: return None`), so no `SchedulerStats` is produced per step — the per-iteration cost is bounded by the existing upstream gate. - The engine core's `Scheduler.make_stats()` also short-circuits inside upstream (`if not self.log_stats: return None`), so no `SchedulerStats` is produced per step — the per-iteration cost is bounded by the existing upstream gate.
@@ -162,10 +175,43 @@ second.
| Metric | Type | Labels | Description | | Metric | Type | Labels | Description |
|--------|------|--------|-------------| |--------|------|--------|-------------|
| `vllm:omni_num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages | | `vllm_omni:num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled | | `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` | | `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds | | `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
### Image and diffusion service-level metrics
| Metric | Type | Labels | Description |
|--------|------|--------|-------------|
| `vllm_omni:stage_gen_time_s` | Histogram | `model_name`, `stage`, `stage_type` | Stage submit to finished output; includes in-stage queueing |
| `vllm_omni:request_queue_wait_s` | Histogram | `model_name` | Orchestration-layer queue wait; a present zero is recorded |
| `vllm_omni:stage_waiting_requests` | Gauge | `model_name`, `stage` | Sum of the latest waiting snapshots across live replicas in the stage |
| `vllm_omni:stage_in_queue_s` | Histogram | `model_name`, `stage` | Diffusion scheduler wait from enqueue to initial execution admission |
| `vllm_omni:num_inference_steps` | Histogram | `model_name` | Denoising-step distribution for diffusion stages |
| `vllm_omni:image_count_total` | Counter | `model_name` | Cumulative number of generated image outputs |
| `vllm_omni:image_pixels` | Histogram | `model_name` | Pixel-count distribution for image outputs |
| `vllm_omni:peak_memory_mb` | Gauge | `model_name`, `stage` | Positive per-stage peak-memory samples reported by engine outputs |
| `vllm_omni:requests_failed_total` | Counter | `model_name`, `reason` | Request failures with bounded reason values (`client_abort`, `client_disconnect`, `stage_error`, `unknown`) |
| `vllm_omni:kv_wait_s` | Histogram | `model_name`, `connector_type` | Completed AR KV-transfer wait intervals, grouped by connector backend |
### Diffusion execution breakdown
Labels: `{model_name, stage, replica}`.
| Metric | Type | Description |
|--------|------|-------------|
| `vllm_omni:diffusion_exec_s` | Histogram | Core diffusion-step execution time per request |
| `vllm_omni:diffusion_exec_per_step_s` | Histogram | Core diffusion execution divided by inference-step count |
| `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time |
| `vllm_omni:diffusion_postprocess_s` | Histogram | Diffusion output postprocessing time |
| `vllm_omni:vae_decode_s` | Histogram | VAE decode time when the engine reports it |
| `vllm_omni:diffusion_forward_s` | Histogram | Forward-only denoise-loop time when pipeline profiling is enabled |
| `vllm_omni:diffusion_kv_load_s` | Histogram | Diffusion-side KV receive/load time when upstream KV is used |
| `vllm_omni:image_ttfp_s` | Histogram | Image-stage submit to first materialized image output |
| `vllm_omni:denoise_step_latency_s` | Histogram | Mean diffusion forward time per denoise step for image output stages |
The optional breakdown and memory families are sparse by design: when an engine does not produce the source measurement, no sample is emitted. This differs from queue timings, where an explicitly measured zero is meaningful and is retained.
### Audio (7) ### Audio (7)
@@ -173,13 +219,13 @@ Labels: `{model_name, stage, replica}` plus the listed extra label.
| Metric | Type | Extra label | Description | | Metric | Type | Extra label | Description |
|--------|------|-------------|-------------| |--------|------|-------------|-------------|
| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame | | `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) | | `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` | | `vllm_omni:audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frames generated | | `vllm_omni:audio_frames_total` | Counter | — | Cumulative audio frames generated |
| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps | | `vllm_omni:audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` | | `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio | | `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
### Cross-stage transfer (4) ### Cross-stage transfer (4)
@@ -187,10 +233,10 @@ Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`.
| Metric | Type | Description | | Metric | Type | Description |
|--------|------|-------------| |--------|------|-------------|
| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes | | `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) | | `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) | | `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) | | `vllm_omni:transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
### LLM stage-level (wrapped `vllm:*`) ### LLM stage-level (wrapped `vllm:*`)
@@ -203,7 +249,7 @@ After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency,
- `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution. - `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution.
- Counters use the `_total` suffix (auto-appended by `prometheus_client`). - Counters use the `_total` suffix (auto-appended by `prometheus_client`).
- Sizes use the `_bytes` suffix. - Sizes use the `_bytes` suffix.
- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm:omni_` / `vllm_omni`. - All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm_omni:` / `vllm_omni`.
## Logging vs. Prometheus ## Logging vs. Prometheus
@@ -66,6 +66,12 @@ the boundary affected by the in-flight stage client/process refactor in
[#5441](https://github.com/vllm-project/vllm-omni/pull/5441). Names and [#5441](https://github.com/vllm-project/vllm-omni/pull/5441). Names and
responsibilities proposed only by that PR are not current contracts. responsibilities proposed only by that PR are not current contracts.
The orchestration loop also has an opt-in event-driven mode
(`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, default off) proposed in
[#5221](https://github.com/vllm-project/vllm-omni/pull/5221). It changes poll
cadence only: the routing, ordering, and terminal-state contracts below hold
identically on both loops.
## Ownership boundary ## Ownership boundary
This document owns `AsyncOmniEngine`, `Orchestrator`, request-state creation, This document owns `AsyncOmniEngine`, `Orchestrator`, request-state creation,
+47
View File
@@ -257,6 +257,53 @@ Lookup checks the marker both before validation and after constructing the
lease so a weaker overlapping lookup cannot return a lease after a stronger lease so a weaker overlapping lookup cannot return a lease after a stronger
validation has denied the artifact. validation has denied the artifact.
If publishing the external deny marker fails, lookup returns a retryable
storage failure and marks the artifact directory with a persistent invalidation
mode before releasing the shared lock. Cooperating readers check that mode both
while opening and immediately before returning a lease. Lookup then makes a
bounded attempt to take the exclusive artifact lock and quarantine the entry;
the quarantine transition removes the internal mode marker. If a cooperating
builder replaces the marked inode while the locks are exchanged, containment
preserves that new publication instead of quarantining it. If neither the
external marker nor the directory-mode marker can be persisted, cross-process
exclusion cannot be guaranteed on an unavailable filesystem and the domain
must remain failed rather than treating the artifact as an ordinary invalid
cache entry.
New publications carry the internal invalidation mode through the staging
rename. Hardening clears it and is the logical authority transition for a
process crash; the temporary- and artifact-parent syncs then make the rename
durable, with an artifact-inode sync first making the hardened mode durable. A
process that dies before hardening leaves a persistently denied
entry, while one that dies after successful hardening leaves an authoritative
ready entry. Any synchronous failure after rename is marked invalid and
contained in quarantine before `PUBLICATION_FAILED` is returned, so that
nonretryable result cannot later resolve as a ready hit. If the filesystem
refuses both the inode marker and the quarantine move, the store instead returns
a retryable storage failure with `outcome_uncertain` details; as with deny-marker
double failure, cross-process exclusion cannot be guaranteed until storage
recovers. Once the artifact-parent sync succeeds, a later artifact- or
build-lock release error is logged and lookup resolves the committed entry
instead of reporting a contradictory publication failure. The same authority
rule applies when a competing publisher already installed byte-identical
semantic content; staging-cleanup errors are retried without downgrading that
validated ready entry to publication failure.
Before widening a hardened artifact for a lifecycle move, the store durably
sets the same invalidation mode. A process crash before rename therefore leaves
the writable source denied. After rename, the store re-hardens the destination
but retains the marker until both parent-directory syncs make the move durable;
only then does it durably clear the marker. An interrupted move therefore leaves
a marked transition that the next exact-key lifecycle operation re-hardens.
Explicit cleanup uses the locked move and retries exact-key
`.cleanup.*` tombstones left by an interrupted removal; a later build performs
the same tombstone reconciliation before producing replacement content.
An operational failure while inspecting a noncooperative competing publication
is a retryable storage failure. The competing entry remains authoritative and
is not mislabeled as corrupt or quarantined merely because its manifest could
not be read transiently.
`inspect_domain()` and `inspect_artifact()` return typed, read-only snapshots. `inspect_domain()` and `inspect_artifact()` return typed, read-only snapshots.
Inspection validation does not create a deny marker. Reported lock activity is Inspection validation does not create a deny marker. Reported lock activity is
point-in-time evidence only; kernel locks are not a persistent owner registry. point-in-time evidence only; kernel locks are not a persistent owner registry.
+1 -1
View File
@@ -1,7 +1,7 @@
# Installation # Installation
!!! important !!! important
vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.26.x requires vLLM 0.26.x; the stable v0.26.0 instructions pin vLLM 0.26.0. vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.28.x requires vLLM 0.28.x; the stable v0.28.0 instructions pin vLLM 0.28.0.
vLLM-Omni supports the following hardware platforms: vLLM-Omni supports the following hardware platforms:
+1 -1
View File
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
### Pre-built wheels ### Pre-built wheels
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source). Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0; the 0.28.0 wheels are published together with the v0.28.0 release this branch produces. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
=== "NVIDIA CUDA" === "NVIDIA CUDA"
@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
vLLM-Omni is built based on vLLM. Please install it with command below. vLLM-Omni is built based on vLLM. Please install it with command below.
```bash ```bash
uv pip install vllm==0.26.0 --torch-backend=auto uv pip install vllm==0.28.0 --torch-backend=auto
``` ```
#### Installation of vLLM-Omni #### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
# --8<-- [start:build-wheel-from-source] # --8<-- [start:build-wheel-from-source]
#### Installation of vLLM #### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library
```bash ```bash
uv pip install vllm==0.26.0 --torch-backend=auto uv pip install vllm==0.28.0 --torch-backend=auto
``` ```
The 0.26.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. The 0.28.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
#### Installation of vLLM-Omni #### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
```bash ```bash
git clone https://github.com/vllm-project/vllm.git git clone https://github.com/vllm-project/vllm.git
cd vllm cd vllm
git checkout v0.26.0 git checkout v0.28.0
``` ```
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file. Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
```bash ```bash
#For CUDA 13.0 (the default for v0.26.0; the wheel filename has no `+cu130` suffix) #For CUDA 13.0 (the default for v0.28.0; the wheel filename has no `+cu130` suffix)
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.26.0/vllm-0.26.0-cp38-abi3-manylinux_2_28_x86_64.whl export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.28.0/vllm-0.28.0-cp38-abi3-manylinux_2_28_x86_64.whl
``` ```
Install vllm with command below (If you have no existing PyTorch). Install vllm with command below (If you have no existing PyTorch).
```bash ```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
--env "HF_TOKEN=$HF_TOKEN" \ --env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \ -p 8091:8091 \
--ipc=host \ --ipc=host \
vllm/vllm-omni:v0.26.0 \ vllm/vllm-omni:v0.28.0 \
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
``` ```
@@ -124,7 +124,7 @@ If you want to specify the base vLLM version:
```bash ```bash
DOCKER_BUILDKIT=1 docker build \ DOCKER_BUILDKIT=1 docker build \
-f docker/Dockerfile.cuda \ -f docker/Dockerfile.cuda \
--build-arg BASE_IMAGE=vllm/vllm-openai:v0.26.0 \ --build-arg BASE_IMAGE=vllm/vllm-openai:v0.28.0 \
-t vllm-omni-cuda . -t vllm-omni-cuda .
``` ```
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
vLLM-Omni is built based on vLLM. Please install it with command below. vLLM-Omni is built based on vLLM. Please install it with command below.
```bash ```bash
uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723 uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
``` ```
#### Installation of vLLM-Omni #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
# --8<-- [start:build-wheel-from-source] # --8<-- [start:build-wheel-from-source]
#### Installation of vLLM #### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library
```bash ```bash
uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723 uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
``` ```
The pre-built 0.26.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. The pre-built 0.28.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
#### Installation of vLLM-Omni #### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
```bash ```bash
git clone https://github.com/vllm-project/vllm.git git clone https://github.com/vllm-project/vllm.git
cd vllm cd vllm
git checkout v0.26.0 git checkout v0.28.0
python3 -m pip install -r requirements/rocm.txt python3 -m pip install -r requirements/rocm.txt
python3 setup.py develop python3 setup.py develop
``` ```
@@ -130,7 +130,7 @@ docker run --rm \
-v ~/.cache/huggingface:/root/.cache/huggingface \ -v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \ --env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \ -p 8091:8091 \
vllm/vllm-omni-rocm:v0.26.0 \ vllm/vllm-omni-rocm:v0.28.0 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
``` ```
@@ -149,7 +149,7 @@ docker run --rm -it \
-v ~/.cache/huggingface:/root/.cache/huggingface \ -v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \ --env "HF_TOKEN=$HF_TOKEN" \
--entrypoint bash \ --entrypoint bash \
vllm/vllm-omni-rocm:v0.26.0 vllm/vllm-omni-rocm:v0.28.0
``` ```
# --8<-- [end:pre-built-images] # --8<-- [end:pre-built-images]
+3 -3
View File
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
source .venv/bin/activate source .venv/bin/activate
# On CUDA # On CUDA
uv pip install vllm==0.26.0 --torch-backend=auto uv pip install vllm==0.28.0 --torch-backend=auto
# On ROCm # On ROCm
uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723 uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
git clone https://github.com/vllm-project/vllm-omni.git git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni cd vllm-omni
@@ -34,7 +34,7 @@ For additional installation methods — please see the [installation guide](inst
!!! note !!! note
It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni. It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.26.0` and vLLM Omni `0.26.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue. If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.28.0` and vLLM Omni `0.28.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
## Offline Inference ## Offline Inference
-1
View File
@@ -37,7 +37,6 @@ GENERAL_EXAMPLE_SLUGS = frozenset(
"text_to_speech", "text_to_speech",
"text_to_video", "text_to_video",
"x_to_text", "x_to_text",
"x_to_video_audio",
} }
) )
+2 -2
View File
@@ -58,7 +58,6 @@ th {
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) | | `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — | | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — | | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -70,6 +69,8 @@ th {
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — | | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — | | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
| `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — | | `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | — | | `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | — |
| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | — | | `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | — |
@@ -94,7 +95,6 @@ th {
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — | | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — | | `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — | | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | | — |
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — | | `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — | | `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — | | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
+45
View File
@@ -811,6 +811,51 @@ If you encounter OOM errors:
Use `/v1/audio/voices` to list available voices for the loaded model. Use `/v1/audio/voices` to list available voices for the loaded model.
## Orchestration Loop (experimental)
Multi-stage omni deployments route stage outputs through a single orchestrator
loop. By default that loop polls every stage replica on a 1 ms cadence. An
opt-in event-driven mode replaces the poll with one reader task per live stage
replica awaiting its client directly, and switches the serving-side
final-output drain to a condition-variable wakeup at the same time.
**Configuration (environment variables):**
| Variable | Default | Description |
|----------|---------|-------------|
| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `0` (off) | Switches the orchestration loop and the final-output drain from the legacy 1 ms poll to event-driven wakeups. Enabled by `1`, `true`, `yes`, or `on`, matched case-insensitively after surrounding whitespace is stripped; any other value leaves it off. |
Set it on the process that runs the orchestrator (stage 0 of an omni
deployment) before starting the server:
```bash
export VLLM_OMNI_EVENT_DRIVEN_ORCH=1
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
--omni \
--port 8091
```
The server logs the selected loop mode and its reader/poller counts once at
startup, so you can confirm which loop is live.
Routing, output ordering, and terminal-state behavior are identical on both
loops; only the poll cadence changes. Leaving the variable unset keeps the
legacy poll loop, which is the supported default.
**Known limitations:**
- The measured serving A/B (idle CPU 2.43% to 0.07%; TTFP p99 -32% at
concurrency 8) predates the rebuild on the per-replica fault-isolation work
in [#4285](https://github.com/vllm-project/vllm-omni/pull/4285). That work
changed dead-replica handling and reader/poller lifecycle rather than the
steady-state output path, and the parity suite covers it, but the serving
A/B has not been re-run on the current head.
- The diffusion-poller branch is covered by unit tests only. Deployments whose
stages all run as standard engine cores never exercise it, including GLM-TTS,
which deploys its DiT without `stage_type: diffusion`.
- Concurrency 1 and 32 measured at parity with the legacy loop. At 32 the
latency is admission-bound, which this mode does not address.
## Development ## Development
Enable debug logging: Enable debug logging:
@@ -112,6 +112,12 @@ perform a collective and uses the same rank-local transfer behavior.
HWR is an opt-in startup optimization for models that declare the final-layout HWR is an opt-in startup optimization for models that declare the final-layout
BF16 restore contract. Use it only with no-AllGather DLO: BF16 restore contract. Use it only with no-AllGather DLO:
The validated BF16 model contracts currently cover MiniMax H3 and
`black-forest-labs/FLUX.2-klein-4B`. FLUX.2-klein-9B shares the same model
class but has not been validated against this contract. FLUX.2-dev, online
FP8, HSDP, LoRA/adapted weights, and non-default load formats remain outside
the validated scope.
```bash ```bash
vllm serve /path/to/model --omni \ vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \ --enable-distributed-layerwise-offload \
+2 -3
View File
@@ -129,7 +129,7 @@ The following tables show which models support each feature:
| **Krea 2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **Krea 2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **OmniGen2** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **OmniGen2** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
@@ -147,7 +147,7 @@ The following tables show which models support each feature:
> Notes: > Notes:
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT. > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary. > 2. `Tongyi-MAI/Z-Image-Turbo` is a distilled model with minimal NFEs; CFG-Parallel is not necessary.
> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload. > 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload.
> 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`. > 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.
> 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution). > 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
@@ -165,7 +165,6 @@ The following tables show which models support each feature:
| **SANA-Video-2B T2V I2V** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | | **SANA-Video-2B T2V I2V** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ |
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* | | **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **LongCat-Video-Avatar-1.5** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **LongCat-Video-Avatar-1.5** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **MiniMax-H3** | ✅ (FL2VA) | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ | | **MiniMax-H3** | ✅ (FL2VA) | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ |
@@ -1,115 +0,0 @@
# X-To-Video-Audio
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/x_to_video_audio>.
The `DreamID-Omni` pipeline generates short videos from text, image and video.
## Local CLI Usage
### Download the Model locally
Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
```bash
python download_dreamid_omni.py --output-dir ./dreamid_omni
```
After download, the model directory will look like this:
```
dreamid_omni/
├── DreamID-Omni/
│ ├── dreamid_omni.safetensors
├── MMAudio/
│ ├── ext_weights/
│ │ ├── best_netG.pt
│ │ ├── v1-16.pth
├── Wan2.2-TI2V-5B/
│ ├── google/*
│ ├── models_t5_umt5-xxl-enc-bf16.pth
│ ├── Wan2.2_VAE.pth
│
├── model_index.json # create by download_dreamid_omni.py
```
### Run the Inference
```python
python x_to_video_audio.py \
--model /path/to/dreamid_omni \
--prompt "Two people walking together and singing happily" \
--image-path ./example0.png ./example1.png \
--audio-path ./example0.wav ./example1.wav \
--video-negative-prompt "jitter, bad hands, blur, distortion" \
--audio-negative-prompt "robotic, muffled, echo, distorted" \
--cfg-parallel-size 2 \
--num-inference-steps 45 \
--height 704 \
--width 1280 \
--output out_dreamid_omni_twoip.mp4
```
In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP:
```python
python x_to_video_audio.py \
--model /path/to/dreamid_omni \
--prompt "Two people walking together and singing happily" \
--image-path ./example0.png ./example1.png \
--audio-path ./example0.wav ./example1.wav \
--video-negative-prompt "jitter, bad hands, blur, distortion" \
--audio-negative-prompt "robotic, muffled, echo, distorted" \
--cfg-parallel-size 2 \
--num-inference-steps 45 \
--height 704 \
--width 1280 \
--use-hsdp \
--hsdp-shard-size 2 \
--output out_dreamid_omni_twoip.mp4
```
You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing.
```python
# Example usage for oneip, ref media from the official repo DreamID-Omni
python x_to_video_audio.py \
--model /path/to/dreamid_omni \
--prompt "<img1>: In the frame, a woman with black long hair is identified as <sub1>.\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: <sub1> is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: <sub1> tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n<sub1> maintains eye contact, swallows as if choosing her words, and says, <S>I keep telling myself I’m fine,but some nights it feels like I’m just performing calm.<E>" \
--image-path 9.png \
--audio-path 9.wav \
--video-negative-prompt "jitter, bad hands, blur, distortion" \
--audio-negative-prompt "robotic, muffled, echo, distorted" \
--cfg-parallel-size 2 \
--num-inference-steps 45 \
--height 704 \
--width 1280 \
--output out_dreamid_omni_oneip.mp4
```
Key arguments:
- `--prompt`: text description (string).
- `--model`: path to the model local directory.
- `--height/--width`: output resolution (defaults 704 * 1024).
- `--image-path`: path to the input image list.
- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks.
- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
- `--hsdp-replicate-size`: number of HSDP replica groups.
- `--num-inference-steps`: number of denoising steps (defaults 45).
- `--video-negative-prompt`: negative prompt for video generation.
- `--audio-negative-prompt`: negative prompt for audio generation.
- `--enable-cpu-offload`: enable CPU offload (defaults False).
## Example materials
??? abstract "download_dreamid_omni.py"
``````py
--8<-- "examples/offline_inference/x_to_video_audio/download_dreamid_omni.py"
``````
??? abstract "x_to_video_audio.py"
``````py
--8<-- "examples/offline_inference/x_to_video_audio/x_to_video_audio.py"
``````
@@ -28,6 +28,18 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config /path/to/deploy_config_file --deploy-config /path/to/deploy_config_file
``` ```
To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct
weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key
`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
```
Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select
the same single-stage pipeline without `--deploy-config`.
### Launch individual stages (stage-based CLI) ### Launch individual stages (stage-based CLI)
Adopt the stage-based CLI architecture to independently instantiate execution processes per functional stage. Adopt the stage-based CLI architecture to independently instantiate execution processes per functional stage.
+50
View File
@@ -0,0 +1,50 @@
# SVDQuant W4A4
## Overview
[SVDQuant](https://arxiv.org/abs/2411.05007) combines four-bit weights
and activations with a small low-rank branch that corrects part of the
quantization error. vLLM-Omni consumes an offline-quantized checkpoint;
it does not calibrate the model while loading.
## Checkpoint contract
Place the following entry in the diffusion transformer's `config.json`:
```json
{
"quantization_config": {
"quant_method": "svdquant",
"rank": 32,
"precision": "nvfp4",
"act_unsigned": false,
"modules_to_not_convert": []
}
}
```
For a quantized linear with input size `K`, output size `N`, and correction
rank `R`, the checkpoint stores:
| Suffix | Shape | dtype |
| --- | --- | --- |
| `qweight` | `(N, K / 2)` | `int8` (two packed FP4 values per byte) |
| `wscales` | `(K / 16, N)` | `float8_e4m3fn` |
| `proj_down` | `(K, R)` | `bfloat16` |
| `proj_up` | `(N, R)` | `bfloat16` |
| `smooth_factor` | `(K,)` | `bfloat16` |
| `wcscales` | `(N,)` | `bfloat16` |
| `wtscale` | `(1,)` | `bfloat16` |
`K` must be divisible by 16 on every tensor-parallel rank. Set `wcscales` to
ones when no per-output correction is needed. Modules listed in
`modules_to_not_convert` keep their checkpoint precision.
## Runtime support
The compatibility path accepts BF16 inputs and executes an NVFP4 GEMM followed
by the BF16 rank correction. It supports vLLM's FlashInfer, CUTLASS, and FBGEMM
NVFP4 tensor layouts; incompatible forced backends fail during model loading.
SM103 is the currently validated and enabled hardware target. Native fusion of
the NVFP4 GEMM and rank correction is separate from this checkpoint-loading
contract.
+9
View File
@@ -0,0 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""vLLM-Omni example scripts, importable as a package.
Keep this ``__init__.py``: with the upstream vLLM installed editable (which also
places an ``examples`` package on ``sys.path``), a regular package at this repo
root wins over the upstream ``examples`` namespace during ``import examples``.
"""
@@ -0,0 +1,117 @@
"""
Offline inference example: MammothModa2 image summarization (single AR stage).
Example:
uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
--model path/to/MammothModa2-Preview \
--deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \
--image /path/to/input.jpg \
--question "Please summarize the content of this image."
"""
from __future__ import annotations
import argparse
import os
from PIL import Image
from vllm import SamplingParams
from vllm.multimodal.image import convert_image_mode
from vllm_omni import Omni
DEFAULT_SYSTEM = "You are a helpful assistant."
DEFAULT_QUESTION = "Please summarize the content of this image."
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.")
parser.add_argument("--image", type=str, required=True, help="Path to input image.")
parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
parser.add_argument(
"--max-tokens",
type=int,
default=512,
help="Max new tokens to generate.",
)
parser.add_argument("--temperature", type=float, default=0.2)
parser.add_argument("--top-p", type=float, default=0.9)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--trust-remote-code", action="store_true")
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
return parser.parse_args()
def build_prompt(system: str, question: str) -> str:
return (
f"<|im_start|>system\n{system}<|im_end|>\n"
"<|im_start|>user\n"
"<|vision_start|><|image_pad|><|vision_end|>"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
def main() -> None:
args = parse_args()
if not os.path.exists(args.image):
raise FileNotFoundError(f"Image file not found: {args.image}")
pil_image = Image.open(args.image)
image_data = convert_image_mode(pil_image, "RGB")
prompt = build_prompt(args.system, args.question)
omni = Omni(
model=args.model,
deploy_config=args.deploy_config,
trust_remote_code=args.trust_remote_code,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
)
try:
sp = SamplingParams(
temperature=float(args.temperature),
top_p=float(args.top_p),
top_k=-1,
max_tokens=int(args.max_tokens),
seed=int(args.seed),
detokenize=True,
)
# NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None].
# Consume it inside the try block so the worker isn't closed early.
outputs = list(
omni.generate(
[
{
"prompt": prompt,
"multi_modal_data": {"image": image_data},
"additional_information": {"omni_task": ["chat"]},
}
],
[sp],
)
)
finally:
omni.close()
lines: list[str] = []
for stage_outputs in outputs:
ro = stage_outputs
text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
lines.append("answer:\n")
lines.append(text.strip() + "\n")
lines.append("\n")
print("\n".join(lines))
if __name__ == "__main__":
main()
@@ -599,6 +599,17 @@ def main():
if args.seed is not None and hasattr(params, "seed"): if args.seed is not None and hasattr(params, "seed"):
params.seed = args.seed params.seed = args.seed
# MammothModa2's AR stage emits one visual token per grid cell,
# one EOL token per row, and one final look-ahead token whose hidden
# state is unavailable. Size the first stage from the prompt metadata
# instead of SamplingParams' default of 16 tokens.
prompt_info = prompt_dict.get("additional_information", {})
if idx == 0 and prompt_info.get("omni_task") == ["t2i"]:
ar_width = int(prompt_info.get("ar_width", [0])[0])
ar_height = int(prompt_info.get("ar_height", [0])[0])
if ar_width > 0 and ar_height > 0:
params.max_tokens = ar_height * (ar_width + 1) + 1
if not diffusion_replaced and len(sampling_params_list) == 1: if not diffusion_replaced and len(sampling_params_list) == 1:
sampling_params_list = [diffusion_params] sampling_params_list = [diffusion_params]
+31 -3
View File
@@ -10,10 +10,12 @@ Generate text from text or image inputs with vLLM-Omni's shared offline entrypoi
| --- | --- | --- | --- | | --- | --- | --- | --- |
| `ByteDance-Seed/BAGEL-7B-MoT` | Yes | Yes | `vllm_omni/deploy/bagel.yaml` | | `ByteDance-Seed/BAGEL-7B-MoT` | Yes | Yes | `vllm_omni/deploy/bagel.yaml` |
| `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | `vllm_omni/deploy/hunyuan_image3_ar.yaml` | | `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | `vllm_omni/deploy/hunyuan_image3_ar.yaml` |
| `bytedance-research/MammothModa2-Preview` | Yes | Yes | `vllm_omni/deploy/mammoth_moda2_ar.yaml` |
| `bytedance-research/MammothModa2-Dev` | Yes | Yes | `vllm_omni/deploy/mammoth_moda2_ar.yaml` |
The script recognizes these model families from `config.json`, applies the The script recognizes these model families from `config.json`, applies the
model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3. model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3
BAGEL uses its registered default deploy. and MammothModa2. BAGEL uses its registered default deploy.
## Text-To-Text ## Text-To-Text
@@ -33,6 +35,22 @@ python examples/offline_inference/x_to_text/x_to_text.py \
--prompt "Explain multimodal inference in three concise sentences." --prompt "Explain multimodal inference in three concise sentences."
``` ```
### MammothModa2
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model bytedance-research/MammothModa2-Preview \
--prompt "Explain multimodal inference in three concise sentences."
```
The same command supports the Dev checkpoint:
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model bytedance-research/MammothModa2-Dev \
--prompt "Explain multimodal inference in three concise sentences."
```
## Image-To-Text ## Image-To-Text
Pass one image with `--image`. The default question can be replaced with any Pass one image with `--image`. The default question can be replaced with any
@@ -54,6 +72,16 @@ python examples/offline_inference/x_to_text/x_to_text.py \
--prompt "Describe this image in detail." --prompt "Describe this image in detail."
``` ```
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model bytedance-research/MammothModa2-Preview \
--image image.png \
--prompt "Describe this image in detail."
```
Use `bytedance-research/MammothModa2-Dev` in the same command to run I2T with
the Dev checkpoint.
## Key Arguments ## Key Arguments
| Argument | Default | Description | | Argument | Default | Description |
@@ -74,4 +102,4 @@ python examples/offline_inference/x_to_text/x_to_text.py \
The committed HunyuanImage-3 AR-only default uses four GPUs with TP=4. To use a The committed HunyuanImage-3 AR-only default uses four GPUs with TP=4. To use a
different topology, pass a custom AR-only YAML through `--deploy-config`. different topology, pass a custom AR-only YAML through `--deploy-config`.
BAGEL uses its repository deploy default unless overridden. BAGEL and MammothModa2 use their repository deploy defaults unless overridden.
@@ -15,6 +15,7 @@ from vllm_omni.model_extras import build_x_to_text_prompt, get_x_to_text_model_f
_REPO_ROOT = Path(__file__).resolve().parents[3] _REPO_ROOT = Path(__file__).resolve().parents[3]
_HUNYUAN_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml" _HUNYUAN_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml"
_MAMMOTH_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "mammoth_moda2_ar.yaml"
def parse_args() -> argparse.Namespace: def parse_args() -> argparse.Namespace:
@@ -61,7 +62,7 @@ def main() -> None:
omni_kwargs: dict[str, Any] = { omni_kwargs: dict[str, Any] = {
"model": args.model, "model": args.model,
"mode": "image-to-text" if image is not None else "text-to-text", "mode": "image-to-text" if image is not None else "text-to-text",
"trust_remote_code": args.trust_remote_code or family == "hunyuan_image3", "trust_remote_code": args.trust_remote_code or family in {"hunyuan_image3", "mammoth_moda2"},
"log_stats": args.log_stats, "log_stats": args.log_stats,
} }
if args.enforce_eager is not None: if args.enforce_eager is not None:
@@ -70,6 +71,8 @@ def main() -> None:
omni_kwargs["deploy_config"] = args.deploy_config omni_kwargs["deploy_config"] = args.deploy_config
elif family == "hunyuan_image3": elif family == "hunyuan_image3":
omni_kwargs["deploy_config"] = str(_HUNYUAN_AR_DEPLOY_CONFIG) omni_kwargs["deploy_config"] = str(_HUNYUAN_AR_DEPLOY_CONFIG)
elif family == "mammoth_moda2":
omni_kwargs["deploy_config"] = str(_MAMMOTH_AR_DEPLOY_CONFIG)
omni = Omni(**omni_kwargs) omni = Omni(**omni_kwargs)
try: try:
@@ -1,108 +0,0 @@
import argparse
import fcntl
import os
import site
import subprocess
import tempfile
import time
from pathlib import Path
from huggingface_hub import snapshot_download
DEPENDENCY_REPO = "https://github.com/bytedance/DreamID-V.git"
DEPENDENCY_BRANCH = "omni"
CACHE_DIR = Path(tempfile.gettempdir()) / "vllm-omni-dependency"
LOCK_FILE = CACHE_DIR / ".install.lock"
DEPENDENCY_DIR = CACHE_DIR / "DreamID-Omni"
def download_dependency():
CACHE_DIR.mkdir(parents=True, exist_ok=True)
with open(LOCK_FILE, "w") as f:
fcntl.flock(f, fcntl.LOCK_EX)
if not DEPENDENCY_DIR.exists():
print(f"Downloading DreamID-Omni to {DEPENDENCY_DIR} ...")
subprocess.run(
["git", "clone", "--depth", "1", DEPENDENCY_REPO, "--branch", DEPENDENCY_BRANCH, str(DEPENDENCY_DIR)],
check=True,
)
print("Download finished.")
fcntl.flock(f, fcntl.LOCK_UN)
# write .pth to site-packages
site_packages = Path(site.getsitepackages()[0])
pth_file = site_packages / "vllm_omni_dependency.pth"
pth_file.write_text(str(DEPENDENCY_DIR))
print(f"Added {DEPENDENCY_DIR} to site-packages via {pth_file}")
def timed_download(repo_id: str, local_dir: str, allow_patterns: list | None = None):
"""Download files from HF repo and log time + destination."""
if os.path.exists(local_dir):
print(f"Directory {local_dir} already exists. Skipping download.")
return
print(f"Starting download from {repo_id} into {local_dir}")
start_time = time.time()
snapshot_download(
repo_id=repo_id,
local_dir=local_dir,
local_dir_use_symlinks=False,
allow_patterns=allow_patterns,
)
elapsed = time.time() - start_time
print(f"✅ Finished downloading {repo_id} in {elapsed:.2f} seconds. Files saved at: {local_dir}")
def main(output_dir: str):
# Wan2.2
wan_dir = os.path.join(output_dir, "Wan2.2-TI2V-5B")
timed_download(
repo_id="Wan-AI/Wan2.2-TI2V-5B",
local_dir=wan_dir,
allow_patterns=["google/*", "models_t5_umt5-xxl-enc-bf16.pth", "Wan2.2_VAE.pth"],
)
# MMAudio
mm_audio_dir = os.path.join(output_dir, "MMAudio")
timed_download(
repo_id="hkchengrex/MMAudio",
local_dir=mm_audio_dir,
allow_patterns=["ext_weights/best_netG.pt", "ext_weights/v1-16.pth"],
)
dreamid_dir = os.path.join(output_dir, "DreamID-Omni")
timed_download(repo_id="XuGuo699/DreamID-Omni", local_dir=dreamid_dir)
# Now we construct the config file
import json
data = {
"_class_name": "DreamIDOmniPipeline",
}
with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
json.dump(data, f, indent=2)
print(f"model_index.json created at {os.path.join(output_dir, 'model_index.json')}")
transformer_dir = os.path.join(output_dir, "transformer")
os.makedirs(transformer_dir, exist_ok=True)
with open(os.path.join(transformer_dir, "config.json"), "w", encoding="utf-8") as f:
json.dump({"fusion": "DreamID-Omni/dreamid_omni.safetensors"}, f)
print(f"transformer/config.json created at {os.path.join(transformer_dir, 'config.json')}")
# now we download the dependency code
download_dependency()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Download models from Hugging Face")
parser.add_argument(
"--output-dir", type=str, default="./dreamid_omni", help="Base directory to save downloaded models"
)
args = parser.parse_args()
main(args.output_dir)
@@ -1,163 +0,0 @@
# X-To-Video-Audio
The `DreamID-Omni` pipeline generates short videos from text, image and video.
## Local CLI Usage
### Download the Model locally
Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
```bash
python download_dreamid_omni.py --output-dir ./dreamid_omni
```
After download, the model directory will look like this:
```
dreamid_omni/
├── DreamID-Omni/
│ ├── dreamid_omni.safetensors
├── MMAudio/
│ ├── ext_weights/
│ │ ├── best_netG.pt
│ │ ├── v1-16.pth
├── Wan2.2-TI2V-5B/
│ ├── google/*
│ ├── models_t5_umt5-xxl-enc-bf16.pth
│ ├── Wan2.2_VAE.pth
│
├── model_index.json
└── transformer/
└── config.json # create by download_dreamid_omni.py
```
### Run the Inference
```python
python x_to_video_audio.py \
--model /path/to/dreamid_omni \
--prompt "Two people walking together and singing happily" \
--image-path ./example0.png ./example1.png \
--audio-path ./example0.wav ./example1.wav \
--video-negative-prompt "jitter, bad hands, blur, distortion" \
--audio-negative-prompt "robotic, muffled, echo, distorted" \
--cfg-parallel-size 2 \
--num-inference-steps 45 \
--height 704 \
--width 1280 \
--output out_dreamid_omni_twoip.mp4
```
In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP:
```python
python x_to_video_audio.py \
--model /path/to/dreamid_omni \
--prompt "Two people walking together and singing happily" \
--image-path ./example0.png ./example1.png \
--audio-path ./example0.wav ./example1.wav \
--video-negative-prompt "jitter, bad hands, blur, distortion" \
--audio-negative-prompt "robotic, muffled, echo, distorted" \
--cfg-parallel-size 2 \
--num-inference-steps 45 \
--height 704 \
--width 1280 \
--use-hsdp \
--hsdp-shard-size 2 \
--output out_dreamid_omni_twoip.mp4
```
You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing.
```python
# Example usage for oneip, ref media from the official repo DreamID-Omni
python x_to_video_audio.py \
--model /path/to/dreamid_omni \
--prompt "<img1>: In the frame, a woman with black long hair is identified as <sub1>.\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: <sub1> is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: <sub1> tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n<sub1> maintains eye contact, swallows as if choosing her words, and says, <S>I keep telling myself I’m fine,but some nights it feels like I’m just performing calm.<E>" \
--image-path 9.png \
--audio-path 9.wav \
--video-negative-prompt "jitter, bad hands, blur, distortion" \
--audio-negative-prompt "robotic, muffled, echo, distorted" \
--cfg-parallel-size 2 \
--num-inference-steps 45 \
--height 704 \
--width 1280 \
--output out_dreamid_omni_oneip.mp4
```
## MagiHuman (`--model-type magi-human`)
MagiHuman is a text → video+audio model with a DiT MoE backbone and a ~9B-param
T5Gemma text encoder. A detailed text prompt is the only required input; an optional
image and/or audio file may be supplied for conditioning. Natively supports Tensor
Parallelism. For an 80GB node, `--tensor-parallel-size 4` is recommended to shard
the MoE weights and the text encoder.
> Install [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for correct
> attention-kernel behaviour (the pipeline otherwise falls back to stubs).
Text-only generation:
```bash
python x_to_video_audio.py \
--model-type magi-human \
--model /path/to/daVinci-MagiHuman \
--prompt "A young woman with long, wavy golden blonde hair... <dialogue and background sound>" \
--tensor-parallel-size 4 \
--height 256 --width 448 \
--num-inference-steps 8 \
--seed 52 \
--extra-body '{"seconds": 5, "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
--output output_magihuman.mp4
```
With optional image and audio conditioning:
```bash
python x_to_video_audio.py \
--model-type magi-human \
--model /path/to/daVinci-MagiHuman \
--prompt "A young woman..." \
--tensor-parallel-size 4 \
--height 256 --width 448 \
--num-inference-steps 8 \
--seed 52 \
--extra-body '{"image_path": "/path/to/ref.jpg", "audio_path": "/path/to/ref.wav", "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
--output output_magihuman.mp4
```
MagiHuman-specific arguments are passed as a JSON dict via `--extra-body` (declared
in `vllm_omni/model_extras/magi_human.py`, routed via `extra_args`):
- `seconds`: output duration in seconds (default 10; ignored when `audio_path` is set,
because the audio length then determines the number of frames).
- `audio_path`: path to an audio file for audio-to-video conditioning. When provided,
the audio drives the frame count and is encoded as a latent condition into the DiT
(`is_a2v` mode). Omit for pure text-to-video+audio generation.
- `image_path`: path to an image file for visual conditioning. Applied at both the
base-resolution (BR) and super-resolution (SR) stages when SR is enabled.
- `sr_height` / `sr_width`: super-resolution output resolution. SR stage is skipped
when these are omitted.
- `sr_num_inference_steps`: denoising steps for the SR stage.
## DreamID-Omni arguments
Key arguments:
- `--prompt`: text description (string).
- `--model`: path to the model local directory.
- `--height/--width`: output resolution (defaults 704 * 1024).
- `--image-path`: path to the input image list.
- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks.
- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
- `--hsdp-replicate-size`: number of HSDP replica groups.
- `--num-inference-steps`: number of denoising steps (defaults 45).
- `--video-negative-prompt`: negative prompt for video generation.
- `--audio-negative-prompt`: negative prompt for audio generation.
- `--enable-cpu-offload`: enable CPU offload (defaults False).
- `--cache-backend`: enable `cache_dit` for acceleration.
- `--quantization`: online (dynamic) quantization method — `fp8` or `int8`.
(VAEs, the T5 text encoder, norms and modulation stay bf16.)
@@ -1,279 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import json
import re
import time
from typing import Any
import numpy as np
from PIL import Image
from vllm.multimodal.media.audio import load_audio
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"}
AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Offline inference for X -> video+audio models.")
parser.add_argument("--model", required=True, help="Model ckpt root directory.")
parser.add_argument(
"--model-type",
default="dreamid-omni",
choices=["dreamid-omni", "magi-human"],
help="Model type.",
)
parser.add_argument("--prompt", default=None, help="Text prompt.")
parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
parser.add_argument("--audio-path", type=str, nargs="+", help="list of audio-path")
parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
parser.add_argument(
"--extra-body",
type=str,
default=None,
help="[magi-human] JSON dict of model-specific extra params (declared in vllm_omni/model_extras/), "
'merged into sampling extra_args. Example: \'{"image_path": "/path/to/img.jpg", "seconds": 5}\'.',
)
parser.add_argument("--height", type=int, default=704, help="Video height.")
parser.add_argument("--width", type=int, default=1280, help="Video width.")
parser.add_argument("--num-inference-steps", type=int, default=45, help="Sampling steps.")
parser.add_argument("--solver-name", default="unipc", help="Solver name: unipc|dpm++|euler.")
parser.add_argument("--shift", type=float, default=5.0, help="Scheduler shift.")
parser.add_argument("--seed", type=int, default=103, help="Random seed for reproducible generation.")
parser.add_argument(
"--cfg-parallel-size",
type=int,
default=1,
choices=[1, 2, 3, 4],
help="Number of GPUs used for classifier free guidance parallel size (max 4 branches).",
)
parser.add_argument(
"--tensor-parallel-size",
type=int,
default=1,
help=("Number of GPUs used for tensor parallelism (TP) inside the DiT."),
)
parser.add_argument(
"--video-negative-prompt",
default="jitter, bad hands, blur, distortion",
help="Negative prompt for video.",
)
parser.add_argument(
"--audio-negative-prompt",
default="robotic, muffled, echo, distorted",
help="Negative prompt for audio.",
)
parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
parser.add_argument(
"--quantization",
type=str,
default=None,
choices=["fp8", "int8"],
help=("Online (dynamic) quantization method for the DreamID-Omni transformer. "),
)
parser.add_argument(
"--enable-cpu-offload",
action="store_true",
default=False,
help="Enable CPU offloading for diffusion models.",
)
parser.add_argument(
"--enable-layerwise-offload",
action="store_true",
help="Enable layerwise (blockwise) offloading on DiT modules.",
)
parser.add_argument("--cache-backend", type=str, default=None, choices=["cache_dit"], help="Cache backend.")
parser.add_argument(
"--use-hsdp",
action="store_true",
help="Enable HSDP for DreamID-Omni fused transformer blocks.",
)
parser.add_argument(
"--hsdp-shard-size",
type=int,
default=1,
help="Number of GPUs used for HSDP sharding when HSDP is enabled.",
)
parser.add_argument(
"--hsdp-replicate-size",
type=int,
default=1,
help="Number of HSDP replica groups. Default 1 means pure sharding.",
)
return parser.parse_args()
def load_image_and_audio(image_paths, audio_paths):
image = []
audio = []
for path in image_paths:
with Image.open(path) as img:
img = img.convert("RGB")
image.append(img)
for path in audio_paths:
audio_array, sr = load_audio(path, sr=16000)
audio_array = audio_array[int(sr * 1) : int(sr * 3)]
audio.append(audio_array)
return image, audio
def _extract_peak_memory_mb(result: Any) -> float:
"""Pull worker-reported peak VRAM (MiB) generation result.
Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb.
"""
if isinstance(result, list):
result = result[0] if result else None
if result is None:
return 0.0
val = getattr(result, "peak_memory_mb", 0.0)
if not val:
inner = result
if isinstance(inner, list):
inner = inner[0] if inner else None
val = getattr(inner, "peak_memory_mb", 0.0)
try:
return float(val or 0.0)
except (TypeError, ValueError):
return 0.0
def main() -> None:
args = parse_args()
if args.prompt is None and args.prompt_file is None:
raise ValueError("Either --prompt or --prompt-file must be provided.")
text_prompt = args.prompt
if args.prompt_file:
with open(args.prompt_file) as f:
text_prompt = json.load(f)
text_prompt = re.sub(
r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", "", text_prompt, flags=re.DOTALL
).strip()
text_prompt = re.sub(
r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]", "", text_prompt, flags=re.DOTALL
).strip()
text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt)
text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip()
if args.model_type == "magi-human":
prompt = text_prompt
extra_args = json.loads(args.extra_body) if args.extra_body else {}
sampling_params = OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
num_inference_steps=args.num_inference_steps,
seed=args.seed,
extra_args=extra_args,
)
else:
image, audio = load_image_and_audio(args.image_path, args.audio_path)
prompt = {
"prompt": text_prompt,
"video_negative_prompt": args.video_negative_prompt,
"audio_negative_prompt": args.audio_negative_prompt,
"multi_modal_data": {"image": image, "audio": audio},
}
sampling_params = OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
num_inference_steps=args.num_inference_steps,
seed=args.seed,
extra_args={
"solver_name": args.solver_name,
"shift": args.shift,
},
)
cache_config = None
if args.cache_backend == "cache_dit":
cache_config = {
"Fn_compute_blocks": 1,
"Bn_compute_blocks": 0,
"max_warmup_steps": 4,
"max_cached_steps": 20,
"residual_diff_threshold": 0.24,
"max_continuous_cached_steps": 3,
"enable_taylorseer": False,
"taylorseer_order": 1,
"scm_steps_mask_policy": None,
"scm_steps_policy": "dynamic",
}
omni_kwargs = dict(
model=args.model,
cfg_parallel_size=args.cfg_parallel_size,
tensor_parallel_size=args.tensor_parallel_size,
use_hsdp=args.use_hsdp,
hsdp_shard_size=args.hsdp_shard_size,
hsdp_replicate_size=args.hsdp_replicate_size,
model_type=args.model_type,
enable_cpu_offload=args.enable_cpu_offload,
enable_layerwise_offload=args.enable_layerwise_offload,
cache_backend=args.cache_backend,
cache_config=cache_config,
)
if args.quantization is not None:
omni_kwargs["quantization"] = args.quantization
omni = Omni(**omni_kwargs)
start = time.perf_counter()
outputs = omni.generate(prompt, sampling_params)
elapsed = time.perf_counter() - start
if not outputs:
raise RuntimeError("No output returned from the model.")
result = outputs[0]
if not result.images:
raise RuntimeError("No video frames found in model output.")
generated_video = result.images[0]
mm = result.multimodal_output or {}
generated_audio = mm.get("audio")
if args.model_type == "magi-human":
# MagiHuman returns frames already as (F, H, W, C) uint8, ready to mux.
fps = float(mm.get("fps", 25))
sample_rate = int(mm.get("audio_sample_rate", 24000))
frames = generated_video
else:
fps = float(mm.get("fps", 24))
sample_rate = int(mm.get("audio_sample_rate", 16000))
# DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1].
# mux_video_audio_bytes expects (F, H, W, C) uint8.
if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4:
raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}")
frames = generated_video.transpose(1, 2, 3, 0)
frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8)
audio_np = None
if generated_audio is not None:
audio_np = np.squeeze(np.asarray(generated_audio)).astype(np.float32)
output_path = args.output
video_bytes = mux_video_audio_bytes(
frames,
audio_np,
fps=float(fps),
audio_sample_rate=sample_rate,
)
with open(output_path, "wb") as f:
f.write(video_bytes)
print(f"Saved generated video to {output_path}")
print(f"Total time: {elapsed:.2f}s")
peak_mb = _extract_peak_memory_mb(outputs)
if peak_mb:
print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
if __name__ == "__main__":
main()
+92 -5
View File
@@ -32,12 +32,14 @@ GPU 0 (90%) and colocates the Talker (55%) and Code2Wav (35%) on GPU 1. That
profile admits at most four concurrent sequences per stage. profile admits at most four concurrent sequences per stage.
| deploy config | GPUs | Notes | | deploy config | GPUs | Notes |
|---|---|---| | --- | --- | --- |
| `minicpmo_4_5.yaml` (default) | 1 | Memory-constrained compatibility layout. | | `minicpmo_4_5.yaml` (default) | 1 | Memory-constrained compatibility layout. |
| `minicpmo_4_5_2gpu.yaml` | 2 | Recommended continuous-batching layout; Talker and Code2Wav share GPU 1. | | `minicpmo_4_5_2gpu.yaml` | 2 | Recommended continuous-batching layout; Talker and Code2Wav share GPU 1. |
| `minicpmo_4_5_3gpu.yaml` | 3 | One GPU per stage. | | `minicpmo_4_5_3gpu.yaml` | 3 | One GPU per stage. |
| `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. | | `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. |
| `minicpmo_4_5_duplex.yaml` | 1 | Experimental native full-duplex overlay. |
Every profile sets `session_mode: duplex`, so native full-duplex is served
from the same process as `/v1/chat/completions`.
The split pipeline preserves native-duplex epoch/turn identity, segment text, The split pipeline preserves native-duplex epoch/turn identity, segment text,
turn completion, reference voice, and terminal-audio metadata through turn completion, reference voice, and terminal-audio metadata through
@@ -55,8 +57,8 @@ vllm-omni serve openbmb/MiniCPM-o-4_5 \
``` ```
For local ModelScope checkpoints, replace `openbmb/MiniCPM-o-4_5` with the For local ModelScope checkpoints, replace `openbmb/MiniCPM-o-4_5` with the
checkpoint path. To start the experimental native duplex backend, use checkpoint path. For native full-duplex, connect `/v1/realtime?duplex=1` on
`vllm_omni/deploy/minicpmo_4_5_duplex.yaml`. this same server.
### Per-stage overrides ### Per-stage overrides
@@ -145,7 +147,7 @@ assistant template, so it is not an apples-to-apples accuracy run.
## Run the Realtime duplex CLI demo ## Run the Realtime duplex CLI demo
After the duplex backend is running, stream one WAV through the Realtime After the server is running, stream one WAV through the Realtime
WebSocket endpoint: WebSocket endpoint:
```bash ```bash
@@ -157,6 +159,33 @@ python examples/online_serving/minicpmo/realtime_duplex_demo.py \
--output-dir /tmp/minicpmo_realtime_duplex_demo --output-dir /tmp/minicpmo_realtime_duplex_demo
``` ```
Video input uses the same session. PyAV demuxes JPEG frames at `--video-fps`
(default 1.0) and vLLM `load_audio` extracts a 16 kHz mono WAV unless
`--input-wav` overrides the soundtrack. Frames keep their capture size
(`--frame-max-side 0`); the server `process_image` normalizes at 448.
`--stack-frames N` raises the visual refresh rate the way official duplex does:
each 1 s unit also samples the `N-1` sub-frames captured inside it and tiles
them into a single composite image sent next to that unit's base frame. The
audio cadence never changes — a unit is always one second — and the wire always
carries 2 images per unit however large `N` is, because the sub-frames share one
composite. Official uses 5 for high refresh rate mode:
```bash
python examples/online_serving/minicpmo/realtime_duplex_demo.py \
--url ws://localhost:8099/v1/realtime?duplex=1 \
--model openbmb/MiniCPM-o-4_5 \
--input-video /path/to/clip.mp4 \
--stack-frames 5 \
--ref-audio /path/to/MiniCPM-o-Demo/assets/ref_audio/ref_minicpm_signature.wav \
--output-dir /tmp/minicpmo_realtime_duplex_video_demo
```
Detail inside a composite is capped by `scale_resolution=448` at
`max_slice_nums=1`: official suggests HD slicing (`max_slice_nums=[2, 1]`) for
stacked frames, which the duplex adapter does not implement yet. Reading small
text or digits out of a wide scene is limited by that, not by frame timing.
## Open the experimental browser client ## Open the experimental browser client
The browser UI serves the page and proxies the same-origin Realtime WebSocket to The browser UI serves the page and proxies the same-origin Realtime WebSocket to
@@ -190,6 +219,64 @@ stronger `response-required` mode is diagnostic: it requires a purpose-built
two-response WAV, its `--input-sha256`, and an two-response WAV, its `--input-sha256`, and an
`--expect-second-response-substring` value. `--expect-second-response-substring` value.
## Run Omni-DuplexEval
Omni-DuplexEval generation uses the vLLM-Omni native MiniCPM-o duplex
endpoint. Evaluation uses a separate OpenAI-compatible multimodal judge
served from the same vLLM-Omni environment. This validates the vLLM-Omni
duplex client and local judge integration; it does not reproduce the paper's
original MiniCPM-o inference implementation.
Start the duplex generation server:
```bash
vllm serve openbmb/MiniCPM-o-4_5 --omni --trust-remote-code \
--deploy-config vllm_omni/deploy/minicpmo_4_5_duplex.yaml \
--served-model-name openbmb/MiniCPM-o-4_5 \
--host 0.0.0.0 --port 8099
```
Start a separate multimodal judge. The allowed path must contain any local
videos passed with `--judge-video-mode video_url`:
```bash
vllm serve Qwen/Qwen2.5-VL-7B-Instruct \
--served-model-name Qwen/Qwen2.5-VL-7B-Instruct \
--allowed-local-media-path /data/omni-duplex-eval \
--host 0.0.0.0 --port 8000
```
Generate, evaluate, and summarize:
```bash
vllm bench omni-duplex-eval --omni generate \
--url ws://127.0.0.1:8099/v1/realtime?duplex=1 \
--model openbmb/MiniCPM-o-4_5 \
--ref-audio /data/ref.wav \
--dataset Hothan/Omni-DuplexEval \
--concurrency 2 \
--response-root /data/omni-duplex-eval/responses
vllm bench omni-duplex-eval --omni evaluate \
--dataset Hothan/Omni-DuplexEval \
--response-root /data/omni-duplex-eval/responses \
--score-root /data/omni-duplex-eval/scores \
--judge-base-url http://127.0.0.1:8000 \
--judge-model Qwen/Qwen2.5-VL-7B-Instruct \
--judge-video-mode video_url \
--eval-workers 4
vllm bench omni-duplex-eval --omni summarize \
--score-root /data/omni-duplex-eval/scores
```
The dataset names such as `RTD_OCR` and `PR_correction` are Hugging Face
splits, not dataset configurations. Pass one with `--split`, or omit it to
run all nine splits. `--limit 1` is useful for a smoke test. Realtime
generation records `clock=media`; artifacts generated with
`--pace as-fast-as-possible` record `clock=invalid` and evaluation rejects
them unless `--allow-invalid-clock` is explicit.
## Related examples ## Related examples
- [Offline MiniCPM-o inference](../../offline_inference/minicpmo/) - [Offline MiniCPM-o inference](../../offline_inference/minicpmo/)
@@ -12,7 +12,9 @@ from __future__ import annotations
import argparse import argparse
import asyncio import asyncio
import base64 import base64
import io
import json import json
import math
import sys import sys
import time import time
from pathlib import Path from pathlib import Path
@@ -22,8 +24,8 @@ if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT)) sys.path.insert(0, str(REPO_ROOT))
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402 from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
PCM16_BYTES_PER_SAMPLE, # noqa: F401 - compatibility export PCM16_BYTES_PER_SAMPLE,
PCM16_SAMPLE_RATE, # noqa: F401 - compatibility export PCM16_SAMPLE_RATE,
RealtimeDuplexClient, RealtimeDuplexClient,
RealtimeEventCollector, RealtimeEventCollector,
build_realtime_url, build_realtime_url,
@@ -38,6 +40,10 @@ from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402 from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
reference_audio_data_url as _ref_audio_data_url, reference_audio_data_url as _ref_audio_data_url,
) )
from vllm_omni.experimental.fullduplex.video_stacking import ( # noqa: E402
concat_frames_b64,
unit_subframe_offsets,
)
class _StreamingOutputWriter: class _StreamingOutputWriter:
@@ -159,14 +165,258 @@ def _event_count_after(
return sum(event.get("type") == event_type for event in events[index + 1 :]) return sum(event.get("type") == event_type for event in events[index + 1 :])
async def run_demo(args: argparse.Namespace) -> dict[str, object]: def _probe_video_frames_and_fps(video_path: Path) -> tuple[int, float]:
input_pcm16 = read_pcm16_wav(Path(args.input_wav)) """Frame count and average FPS via PyAV (same approach as Daily-Omni)."""
if not input_pcm16: import av
raise ValueError("input WAV has no audio")
with av.open(str(video_path)) as container:
stream = container.streams.video[0]
avg_fps = float(stream.average_rate) if stream.average_rate else 0.0
num_frames = stream.frames or 0
if num_frames <= 0:
duration: float | None = None
if stream.duration is not None and stream.time_base is not None:
duration = float(stream.duration * stream.time_base)
elif container.duration is not None:
duration = container.duration / av.time_base
if duration and avg_fps:
num_frames = int(math.floor(duration * avg_fps))
if num_frames <= 0:
stream.thread_type = "AUTO"
num_frames = sum(1 for _ in container.decode(stream))
if num_frames <= 0:
raise ValueError(f"No decodable video frames in {video_path}")
return num_frames, avg_fps or 1.0
def _upright(image, rotation: int):
"""Apply a container display-matrix rotation, in degrees, to ``image``.
PyAV hands back the stored raster and leaves the display matrix for the
caller, unlike the ffmpeg CLI which autorotates. Phone clips routinely
carry +/-90, so skipping this feeds the model a sideways frame: arrows and
digits stop being readable and the model mirrors the motion it reports.
"""
from PIL import Image
quarter_turns = {
90: Image.Transpose.ROTATE_90,
180: Image.Transpose.ROTATE_180,
270: Image.Transpose.ROTATE_270,
}
turn = quarter_turns.get(int(rotation) % 360)
return image.transpose(turn) if turn is not None else image
def _decode_video_frames_rgb(video_path: Path, frame_idx: list[int]) -> list:
"""Decode presentation-order ``frame_idx`` into upright RGB ``PIL.Image``s."""
import av
wanted = sorted({int(i) for i in frame_idx})
if not wanted:
return []
decoded: dict[int, object] = {}
with av.open(str(video_path)) as container:
stream = container.streams.video[0]
stream.thread_type = "AUTO"
cursor = 0
for position, frame in enumerate(container.decode(stream)):
if wanted[cursor] > position:
continue
image = _upright(frame.to_image(), getattr(frame, "rotation", 0) or 0)
while cursor < len(wanted) and wanted[cursor] <= position:
decoded[wanted[cursor]] = image
cursor += 1
if cursor >= len(wanted):
break
if not decoded:
raise ValueError(f"Decoded no frames from {video_path} for indices {wanted[:8]}")
last_image = decoded[max(decoded)]
return [decoded.get(int(i), last_image) for i in frame_idx]
def _resize_max_side(image, max_side: int):
width, height = image.size
longest = max(width, height)
if longest <= max_side or max_side <= 0:
return image
scale = max_side / float(longest)
new_size = (max(1, int(round(width * scale))), max(1, int(round(height * scale))))
return image.resize(new_size)
def _jpeg_bytes(image, *, quality: int = 95) -> bytes:
buf = io.BytesIO()
image.convert("RGB").save(buf, format="JPEG", quality=quality)
return buf.getvalue()
def _to_pcm16_bytes(audio_np) -> bytes:
import numpy as np
arr = np.asarray(audio_np).reshape(-1)
if arr.dtype != np.int16:
arr = (np.clip(arr, -1.0, 1.0) * 32767.0).astype(np.int16)
return arr.tobytes()
def _pcm16_wav_duration_s(path: Path) -> float:
pcm16 = read_pcm16_wav(path)
return len(pcm16) / float(PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE)
def _extract_video_frames(
video_path: Path,
work_dir: Path,
*,
fps: float,
max_side: int,
duration_s: float,
stack_frames: int = 1,
) -> tuple[list[str], list[str | None]]:
"""Demux JPEG frames only. Frame count follows ``duration_s``, not soundtrack.
Returns the base frames (one per model unit) and the parallel stacked track.
With ``stack_frames=N`` each unit also samples its ``N-1`` interior
sub-frames and tiles them into one composite, which is how official duplex
raises the visual refresh rate while keeping one second of audio per unit.
A unit whose interior falls past the end of the clip stacks nothing.
"""
num_frames, avg_fps = _probe_video_frames_and_fps(video_path)
sample_fps = max(float(fps), 1e-6)
unit_s = 1.0 / sample_fps
timeline_s = max(float(duration_s), unit_s)
num_samples = max(1, int(math.ceil(timeline_s * sample_fps - 1e-9)))
offsets = unit_subframe_offsets(stack_frames, unit_s=unit_s)
def frame_at(seconds: float) -> int:
# Official get_video_frame_audio_segments uses truncating int(), not
# round(), so 29.97 fps lands on the same presentation index.
return min(int(seconds * avg_fps), num_frames - 1)
base_idx = [frame_at(unit * unit_s) for unit in range(num_samples)]
sub_times = [
[unit * unit_s + offset for offset in offsets if unit * unit_s + offset < timeline_s]
for unit in range(num_samples)
]
sub_idx = [frame_at(seconds) for times in sub_times for seconds in times]
images = _decode_video_frames_rgb(video_path, base_idx + sub_idx)
sub_images = iter(images[len(base_idx) :])
frame_dir = work_dir / "frames"
frame_dir.mkdir(parents=True, exist_ok=True)
for stale in frame_dir.glob("*.jpg"):
stale.unlink()
frames_b64: list[str] = []
stacked_b64: list[str | None] = []
for index, (image, times) in enumerate(zip(images[: len(base_idx)], sub_times), start=1):
jpeg = _jpeg_bytes(_resize_max_side(image, max_side))
(frame_dir / f"frame_{index:05d}.jpg").write_bytes(jpeg)
frames_b64.append(base64.b64encode(jpeg).decode("ascii"))
interior = [_resize_max_side(next(sub_images), max_side) for _ in times]
if not interior:
stacked_b64.append(None)
continue
composite = concat_frames_b64(interior)
(frame_dir / f"stack_{index:05d}.jpg").write_bytes(base64.b64decode(composite))
stacked_b64.append(composite)
return frames_b64, stacked_b64
def _extract_video_soundtrack(video_path: Path, work_dir: Path) -> Path:
"""Write a 16 kHz mono WAV from the video soundtrack."""
try:
from vllm.multimodal.media.audio import load_audio
except ImportError as exc: # pragma: no cover
raise SystemExit("--input-video needs vLLM audio helpers (load_audio)") from exc
audio_np, sr = load_audio(str(video_path), sr=PCM16_SAMPLE_RATE, mono=True)
work_dir.mkdir(parents=True, exist_ok=True)
audio_path = work_dir / "input_16k.wav"
write_pcm16_wav(audio_path, _to_pcm16_bytes(audio_np), sample_rate_hz=int(sr))
return audio_path
def _extract_video_input(
video_path: Path,
work_dir: Path,
*,
fps: float,
max_side: int,
stack_frames: int,
) -> tuple[Path, list[str], list[str | None]]:
"""Demux a video into duplex inputs: 16 kHz mono WAV plus JPEG frames.
Use this only when the video soundtrack is the audio source. When the
caller already has ``--input-wav``, extract frames against that WAV
duration instead so a silent clip cannot fail ``load_audio``.
"""
audio_path = _extract_video_soundtrack(video_path, work_dir)
frames_b64, stacked_b64 = _extract_video_frames(
video_path,
work_dir,
fps=fps,
max_side=max_side,
duration_s=_pcm16_wav_duration_s(audio_path),
stack_frames=stack_frames,
)
return audio_path, frames_b64, stacked_b64
def _resolve_duplex_av_inputs(
*,
input_wav: str | None,
input_video: str | None,
work_dir: Path,
fps: float,
max_side: int,
stack_frames: int = 1,
) -> tuple[str, list[str], list[str | None]]:
"""Return ``(wav_path, frames_b64, stacked_b64)``. External WAV wins over video audio."""
if input_video:
video_path = Path(input_video)
if input_wav:
frames, stacked = _extract_video_frames(
video_path,
work_dir,
fps=fps,
max_side=max_side,
duration_s=_pcm16_wav_duration_s(Path(input_wav)),
stack_frames=stack_frames,
)
return input_wav, frames, stacked
extracted_wav, frames, stacked = _extract_video_input(
video_path,
work_dir,
fps=fps,
max_side=max_side,
stack_frames=stack_frames,
)
return str(extracted_wav), frames, stacked
if not input_wav:
raise SystemExit("provide --input-wav, --input-video, or both")
return input_wav, [], []
async def run_demo(args: argparse.Namespace) -> dict[str, object]:
output_dir = Path(args.output_dir) output_dir = Path(args.output_dir)
stream_writer = _StreamingOutputWriter(output_dir) stream_writer = _StreamingOutputWriter(output_dir)
input_wav, video_frames, stacked_video_frames = _resolve_duplex_av_inputs(
input_wav=args.input_wav,
input_video=args.input_video,
work_dir=output_dir / "video_input",
fps=args.video_fps,
max_side=args.frame_max_side,
stack_frames=args.stack_frames,
)
input_pcm16 = read_pcm16_wav(Path(input_wav))
if not input_pcm16:
raise ValueError("input WAV has no audio")
url = build_realtime_url( url = build_realtime_url(
args.url, args.url,
args.model, args.model,
@@ -184,10 +434,12 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
timeout_s=args.timeout_s, timeout_s=args.timeout_s,
) )
stream_event_cursor = len(client.events.events) stream_event_cursor = len(client.events.events)
await client.stream_pcm16( frames_sent = await client.stream_pcm16(
input_pcm16, input_pcm16,
chunk_ms=args.chunk_ms, chunk_ms=args.chunk_ms,
realtime=not args.no_realtime_pacing, realtime=not args.no_realtime_pacing,
video_frames=video_frames,
stacked_video_frames=stacked_video_frames,
) )
commit_event_cursor = len(client.events.events) commit_event_cursor = len(client.events.events)
stream_decision = _latest_model_decision(client.events.events, stream_event_cursor) stream_decision = _latest_model_decision(client.events.events, stream_event_cursor)
@@ -310,6 +562,10 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
committed_index, committed_index,
), ),
}, },
"input_video": args.input_video,
"video_frames_extracted": len(video_frames),
"video_frames_sent": frames_sent,
"video_stacked_frames_extracted": sum(frame is not None for frame in stacked_video_frames),
"audio_bytes": len(audio), "audio_bytes": len(audio),
"audio_chunk_count": len(stream_writer.audio_chunk_paths), "audio_chunk_count": len(stream_writer.audio_chunk_paths),
"audio_chunk_files": [str(path) for path in stream_writer.audio_chunk_paths], "audio_chunk_files": [str(path) for path in stream_writer.audio_chunk_paths],
@@ -356,7 +612,48 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1") parser.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1")
parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5") parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
parser.add_argument("--session-id") parser.add_argument("--session-id")
parser.add_argument("--input-wav", required=True) parser.add_argument(
"--input-wav",
help="Mono 16 kHz PCM16 clip to stream. Optional when --input-video carries the audio track.",
)
parser.add_argument(
"--input-video",
help=(
"Video to stream as omni duplex input. PyAV demuxes frames and "
"vLLM load_audio extracts a 16 kHz mono WAV (unless --input-wav overrides the audio)."
),
)
parser.add_argument(
"--video-fps",
type=float,
default=1.0,
help="Frame extraction rate. The wire contract carries ~1 frame per second of audio.",
)
parser.add_argument(
"--frame-max-side",
type=int,
default=0,
help=(
"Optional client-side downscale of the longest side, in pixels. "
"0 (default) keeps the capture resolution so the server "
"process_image can normalize at scale_resolution=448 from a "
"higher-quality source. The web demo never resizes on the client."
),
)
parser.add_argument(
"--stack-frames",
type=int,
default=1,
help=(
"Frames sampled per 1 s audio unit. 1 (default) sends the unit's "
"frame alone. N > 1 also samples the N-1 sub-frames inside the unit "
"and tiles them into one composite sent next to the base frame, so "
"motion (arrows, doors, digits changing) is visible at N fps while "
"the audio stays 1 s per unit. Official uses 5 for high refresh "
"rate mode; the wire carries 2 images per unit at any N because the "
"sub-frames share one composite."
),
)
parser.add_argument( parser.add_argument(
"--ref-audio", "--ref-audio",
required=True, required=True,
@@ -372,7 +669,7 @@ def parse_args() -> argparse.Namespace:
"--temperature", "--temperature",
type=float, type=float,
default=None, default=None,
help="Stage0 sampling temperature; response-required soft-interrupt uses 0.0.", help="Stage 0 sampling temperature; omit to preserve the model default.",
) )
parser.add_argument("--no-realtime-pacing", action="store_true") parser.add_argument("--no-realtime-pacing", action="store_true")
parser.add_argument("--require-audio", action="store_true") parser.add_argument("--require-audio", action="store_true")
@@ -24,6 +24,18 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config /path/to/your_deploy_config.yaml --deploy-config /path/to/your_deploy_config.yaml
``` ```
To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct
weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key
`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
```
Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select
the same single-stage pipeline without `--deploy-config`.
For a 3x-GPU multi-replica layout (talker/code2wav scale-out on cuda:1,2), For a 3x-GPU multi-replica layout (talker/code2wav scale-out on cuda:1,2),
use `--stage-overrides` on top of the default config: use `--stage-overrides` on top of the default config:
+5 -1
View File
@@ -71,11 +71,11 @@ dev = [
"voxcpm==2.0.3", # for VoxCPM2 path-locating in tests "voxcpm==2.0.3", # for VoxCPM2 path-locating in tests
"torchdiffeq==0.2.5", # for covo_audio Test "torchdiffeq==0.2.5", # for covo_audio Test
"FlagEmbedding==1.4.0", # for hunyuan_image3 tests "FlagEmbedding==1.4.0", # for hunyuan_image3 tests
"kernels==0.14.1", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
"s3tokenizer==0.3.0", # for tts test "s3tokenizer==0.3.0", # for tts test
# Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo, # Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo,
# whose top-level stepaudio2 package conflicts with this distribution. # whose top-level stepaudio2 package conflicts with this distribution.
"step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml) "step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml)
"silero-vad==6.1", # Realtime server-VAD duplex E2E
] ]
demo = [ demo = [
@@ -85,6 +85,10 @@ demo = [
"cn2an>=0.5.22", # Chinese number normalization for TTS (e.g. "1次" -> "一次") "cn2an>=0.5.22", # Chinese number normalization for TTS (e.g. "1次" -> "一次")
] ]
server-vad = [
"silero-vad==6.1",
]
forced-aligner = [ forced-aligner = [
"qwen-asr>=0.0.6", # faithful Japanese/Korean word segmentation for the TTS forced aligner "qwen-asr>=0.0.6", # faithful Japanese/Korean word segmentation for the TTS forced aligner
] ]
+47 -5
View File
@@ -10,14 +10,22 @@ checkpoint for T2V and first-frame I2V. Every output includes synchronized
48 kHz stereo audio. 48 kHz stereo audio.
The raw [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5) The raw [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5)
repository is not directly loadable with `--model`. It supplies the official repository is not directly loadable with `--model`. It supplies the canonical
upsampler and LoRA sidecars used by the Full/SFT two-stage pipeline, so accept Diffusion VAE checkpoint and the official upsampler and LoRA sidecars used by
both model licenses and authenticate before first use. the Full/SFT two-stage pipeline, so accept both model licenses and authenticate
before first use.
The vLLM-Omni integration code is provided under the project's
[Apache-2.0 license](../../LICENSE). The LTX-2.5 model weights are separate
artifacts covered by the
[LTX-2.x Community License](https://github.com/Lightricks/LTX-2/blob/main/LICENSE.md),
not the vLLM-Omni license. Users are responsible for reviewing and complying
with the model license before downloading or using the weights.
## Choose a pipeline ## Choose a pipeline
| Pipeline | Mode | Default size | Schedule | | Pipeline | Mode | Default size | Schedule |
|---|---|---:|---:| | --- | --- | ---: | ---: |
| `LTX2Pipeline` | Full/SFT one-stage | 960x544 | 30 steps | | `LTX2Pipeline` | Full/SFT one-stage | 960x544 | 30 steps |
| `LTX2TwoStagePipeline` | Full/SFT two-stage | 1920x1088 | 30 + 3 steps | | `LTX2TwoStagePipeline` | Full/SFT two-stage | 1920x1088 | 30 + 3 steps |
| `LTX2DistilledOneStagePipeline` | Distilled one-stage | 960x544 | 8 steps | | `LTX2DistilledOneStagePipeline` | Distilled one-stage | 960x544 | 8 steps |
@@ -28,6 +36,40 @@ latent upsampler, and run a three-step refinement stage. Select the class with
`--model-class-name`; no `--task-type` flag is required. Supplying one initial `--model-class-name`; no `--task-type` flag is required. Supplying one initial
image selects I2V, while omitting it selects T2V. image selects I2V, while omitting it selects T2V.
## Video decoder
LTX-2.5 uses the canonical Native Diffusion VAE decoder by default. To opt in
to the legacy convolutional VAE, set the startup-only stage-0 model extra
`ltx2_use_conv_vae: true`. The choice applies to all requests and all four
pipeline classes above. DiffVAE loads its NATTEN kernel from Hugging Face Hub
during startup. `decode_timestep` and `decode_noise_scale` condition only the
legacy ConvVAE and have no effect when DiffVAE is selected.
For offline Python usage:
```python
omni = Omni(
model="Lightricks/LTX-2.5-Diffusers",
model_class_name="LTX2Pipeline",
stage_overrides='{"0":{"extras":{"ltx2_use_conv_vae":true}}}',
)
```
For online serving:
```bash
vllm serve Lightricks/LTX-2.5-Diffusers \
--omni \
--model-class-name LTX2Pipeline \
--stage-overrides '{"0":{"extras":{"ltx2_use_conv_vae":true}}}'
```
Both decoders are untiled by default. Set `vae_use_tiling` for memory-saving
serial tiling; DiffVAE tiles only above 80 frames or 768 pixels in either spatial
dimension. DiffVAE also supports distributed VAE decode; see the
[VAE Parallelism Guide](../../docs/user_guide/diffusion/parallelism/vae_parallelism.md).
DiffVAE is decoder-only, so I2V still uses the convolutional VAE encoder.
## Prerequisites ## Prerequisites
```bash ```bash
@@ -42,7 +84,7 @@ Install matching vLLM and vLLM-Omni versions, and ensure `ffmpeg` and
## Hardware ## Hardware
| GPU | Status | Recommended scope | | GPU | Status | Recommended scope |
|---|---|---| | --- | --- | --- |
| NVIDIA B300 | Verified | All four canonical pipelines | | NVIDIA B300 | Verified | All four canonical pipelines |
| NVIDIA B200 or H200 | Capacity-based recommendation; not yet verified | All four canonical pipelines | | NVIDIA B200 or H200 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
| NVIDIA GB200 or GB300 | Capacity-based recommendation; not yet verified | All four canonical pipelines | | NVIDIA GB200 or GB300 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
+212
View File
@@ -0,0 +1,212 @@
# MammothModa2
> MammothModa2-Preview and MammothModa2-Dev unified understanding and generation
## Summary
- Vendor: ByteDance Research
- Models: `bytedance-research/MammothModa2-Preview`, `bytedance-research/MammothModa2-Dev`
- Tasks: Preview and Dev text-to-image (AR → DiT); Dev text/image understanding
- Mode: Offline inference
- Maintainer: Community
## When to use this recipe
Use this recipe to run MammothModa2-Preview through shared task-oriented
offline examples. Text-to-image uses the shared image example (`text_to_image.py`) instead of a model-specific script.
The generic example formats the AR prompt, drives the AR → DiT stage pipeline,
and forwards MammothModa2-specific generation parameters through the
pipeline-declared `extra_body` contract.
MammothModa2's DiT stage consumes its inputs through the multi-stage kwargs
interface (not `OmniDiffusionRequest`), so its generation knobs
(`text_guidance_scale`, `cfg_range`, `num_inference_steps`) are passed via
`--extra-body` rather than the standard `--num-inference-steps` / `--cfg-scale`
flags. Image size uses the standard `--height` / `--width` flags.
## References
- Upstream model:
[`bytedance-research/MammothModa2-Preview`](https://huggingface.co/bytedance-research/MammothModa2-Preview)
- Dev model:
[`bytedance-research/MammothModa2-Dev`](https://huggingface.co/bytedance-research/MammothModa2-Dev)
- Related offline example:
[`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
- Related T2T/I2T example:
[`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.py)
- Declared parameters:
[`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)
- Deploy config:
[`vllm_omni/deploy/mammoth_moda2.yaml`](../../vllm_omni/deploy/mammoth_moda2.yaml)
## Hardware Support
The default deploy config runs both the AR and DiT stages on a single GPU
(`devices: "0"`). The committed `gpu_memory_utilization` split (stage-0 AR `0.5`,
stage-1 DiT `0.3`) is sized for an ~80 GB GPU. The model also fits on a 48 GB GPU
after rebalancing the split so the AR weights (~23 GB) leave room for the KV
cache — see the note under *1x L40S 48GB*.
## GPU
### 1x L40S 48GB
> **48 GB config adjustment:** the committed
> `vllm_omni/deploy/mammoth_moda2.yaml` uses
> `gpu_memory_utilization` 0.5 / 0.3 (sized for ~80 GB). To fit on a 48 GB L40S,
> set the stage-0 (AR) value to `0.8` and the stage-1 (DiT) value to `0.16`
> before running. (On an ~80 GB GPU, leave the defaults unchanged.)
#### Environment
- OS: Linux
- Python: Match the repository requirements for your checkout
- Driver / runtime: NVIDIA CUDA environment with one L40S 48 GB (verified) or an
~80 GB GPU for the default config
- vLLM version: Match the repository requirements for your checkout
- vLLM-Omni version or commit: Use the commit you are deploying from
#### Offline Commands
Download the model:
```bash
hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview
```
Run text-to-image with the shared offline example from the repository root. The
deploy config sets `trust_remote_code`, so no extra flag is needed. Forward the
MammothModa2 generation parameters as a JSON object through `--extra-body`:
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model ./MammothModa2-Preview \
--deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
--height 1024 \
--width 1024 \
--extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
--output mammoth_t2i.png
```
The `--extra-body` JSON forwards MammothModa2-specific parameters into
`OmniDiffusionSamplingParams.extra_args`. Keys are filtered against the model's
declared `extra_body_params` (see
[`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)),
so unknown keys for MammothModa2 are silently dropped:
- `text_guidance_scale` — classifier-free guidance scale for the DiT stage
(default `9.0`; CFG is active only when `> 1.0`).
- `cfg_range` — relative step range `[start, end]` over which CFG is applied
(default `[0.0, 1.0]`).
- `num_inference_steps` — number of DiT denoising steps (default `50`).
`--height` and `--width` must be multiples of 16.
Run text-to-text through the shared understanding example. It recognizes the
MammothModa2 checkpoint and automatically selects `mammoth_moda2_ar.yaml`:
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model ./MammothModa2-Preview \
--prompt "Explain multimodal generation in three sentences."
```
Add an image for image-to-text or image summarization. The shared example
uses MammothModa2's chat and vision-token template:
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model ./MammothModa2-Preview \
--image /path/to/input.jpg \
--prompt "Please summarize the content of this image."
```
#### Verification
The example writes the generated image to the `--output` path. Confirm the file
exists and is a valid image:
```bash
ls -lh mammoth_t2i.png
python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)"
```
### 1x AMD MI300X, MammothModa2 Preview
#### Environment
- OS: Linux 6.8.0-134-generic, x86_64
- Container: official ROCm image built from `docker/Dockerfile.rocm`
- Python: 3.12.13
- PyTorch: 2.11.0+gitd0c8b1f
- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
- vLLM version: 0.27.0+rocm723
- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
#### Offline Commands
The checked run used the committed stage split, with `gpu_memory_utilization` set to 0.5 for AR and 0.3 for DiT:
```bash
python3 examples/offline_inference/text_to_image/text_to_image.py \
--model bytedance-research/MammothModa2-Preview \
--deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
--height 1024 \
--width 1024 \
--seed 42 \
--extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
--enable-diffusion-pipeline-profiler \
--log-stats \
--output mammoth_t2i.png
```
#### Verification
The first request took 85.224 seconds. The AR stage generated 4,161 visual tokens in 72.996 seconds, and the DiT stage took 12.163 seconds. AR weight loading used 21.4 GiB and took 8.250 seconds. DiT weight loading used 5.49 GiB and took 1.824 seconds. The largest one second whole device memory sample was 106.57 GiB, including the AR KV cache reserved by the 0.5 memory setting.
The output was a valid 1024 by 1024 RGB PNG.
## MammothModa2-Dev unified inference
MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses
Qwen2.5-VL. vLLM-Omni selects the matching implementation from the nested
`llm_config.model_type`; no checkpoint edits or `trust_remote_code` flag are
required.
Text-to-text and image-to-text use the AR-only deploy. Text-to-image loads the
Qwen3 generation experts (`gen_mlp`), extra visual vocabulary and image head,
then sends the generated visual tokens and hidden states to the DiT stage.
Download the checkpoint:
```bash
hf download bytedance-research/MammothModa2-Dev --local-dir ./MammothModa2-Dev
```
Run text-to-text through the shared understanding example. It recognizes the
Dev checkpoint as MammothModa2 and automatically selects
`mammoth_moda2_ar.yaml`:
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model ./MammothModa2-Dev \
--prompt "Explain multimodal generation in three sentences."
```
Add an image for image-to-text or image summarization:
```bash
python examples/offline_inference/x_to_text/x_to_text.py \
--model ./MammothModa2-Dev \
--image ./image.png \
--prompt "Please summarize the content of this image."
```
The Dev checkpoint is approximately 47.55 GiB on disk. In the verified AR-only
run, loaded model weights used approximately 16.97 GiB of GPU memory before KV
and encoder caches. Allow additional GPU memory for those caches and the input
image.
@@ -0,0 +1,82 @@
# MiniMax H3 disaggregated text encoder
This opt-in topology runs the Qwen3-VL text encoder as a vLLM stage and sends
its hidden states and token-role metadata to an encoder-free diffusion stage.
The standard MiniMax H3 recipes remain single-stage and continue to load the
text encoder inside the diffusion pipeline.
## Start the server
Choose the topology explicitly and load its deployment defaults:
```bash
vllm-omni serve MiniMaxAI/MiniMax-H3 \
--omni \
--deploy-config vllm_omni/deploy/minimax_h3_disaggregated.yaml
```
The default deployment assigns stage 0 to GPUs 0-1 with tensor parallel size
2 and `max_num_seqs: 1`. Stage 1 uses GPUs 2-5 with diffusion tensor parallel
size 1, Ulysses degree 4, and VAE patch parallel size 4. Adjust the
`devices`, `tensor_parallel_size`, and stage 1 `parallel_config` values in a
deployment override for the available hardware. Diffusion quantization,
layerwise offload, distributed layerwise offload, VAE parallelism, and USP
settings use the same stage 1 options documented in [MiniMax-H3.md](MiniMax-H3.md).
For example, this five-GPU topology assigns one GPU to the encoder and four
to the diffusion stage. `--stage-overrides` keeps placement and parallelism
scoped to the owning stage rather than broadcasting an override to both:
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3,4 \
vllm-omni serve MiniMaxAI/MiniMax-H3 \
--omni \
--deploy-config vllm_omni/deploy/minimax_h3_disaggregated.yaml \
--stage-overrides '{"0":{"devices":"0","tensor_parallel_size":1},"1":{"devices":"1,2,3,4","tensor_parallel_size":1,"ulysses_degree":4,"vae_patch_parallel_size":4}}'
```
For memory-constrained deployments, start from the CPU-offload or distributed
layerwise-offload profiles in [MiniMax-H3.md](MiniMax-H3.md). Apply memory and
quantization options only to Stage 1 with `--stage-overrides`; retain the
encoder's BF16 configuration and the video/audio VAEs' FP32 precision. Select
one offload strategy per deployment:
```bash
# Stage 1 model-level CPU offload.
--stage-overrides '{"1":{"enable_cpu_offload":true}}'
# Stage 1 distributed layerwise offload. Tune resident layers for available RAM.
--stage-overrides '{"1":{"enable_distributed_layerwise_offload":true,"dlo_use_allgather":false,"dlo_resident_layers":20}}'
# Stage 1 online FP8 quantization of the DiT only.
--stage-overrides '{"1":{"diffusion_quantization_config":"{\"transformer\":{\"method\":\"fp8\"}}"}}'
```
The Stage 1 VAE patch-parallel options remain independent of offload and
quantization. See [MiniMax-H3.md](MiniMax-H3.md) for memory requirements and
hardware-qualified profiles before combining these options.
Stage 1 sets `model_loaded.text_encoder: false`; it must not load or download
text-encoder weights. This H3 topology explicitly keeps its single-replica
diffusion stage inline, avoiding serialization of decoded video through a
subprocess. It expects both stages to run in one deployment; cross-node payload
transport is outside this configuration.
The `/v1/videos` request schema and `extra_params.task` values (`t2va`,
`fl2va`, and `ref2va`) are unchanged from the single-stage recipe.
## Turbo LoRA
MiniMax-H3 Turbo uses five sigma points, `flow_shift=6`, and
`audio_flow_shift=3`. Start Turbo deployments with the dedicated defaults so
requests that omit sampling controls do not inherit the 50-step base schedule:
```bash
vllm-omni serve MiniMaxAI/MiniMax-H3 \
--omni \
--lora-path /path/to/MiniMax-H3-Turbo \
--deploy-config vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml
```
The base deployment intentionally retains 50 inference steps for non-LoRA
quality. Turbo LoRA supports T2VA and FL2VA requests, not Ref2VA.
+5 -4
View File
@@ -839,7 +839,7 @@ hf download lightx2v/Minimax-h3-Turbo "${TURBO_FILE}" --local-dir "${TURBO_DIR}"
export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}" export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}"
``` ```
Start from a non-offloaded FL2VA server command and add Start from a non-offloaded or DLO FL2VA server command and add
`--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`. `--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`.
`--lora-path` preloads the adapter; each request still activates it and uses `--lora-path` preloads the adapter; each request still activates it and uses
the published sampling settings: the published sampling settings:
@@ -853,9 +853,10 @@ the published sampling settings:
For FL2VA, change `task` and add `input_reference` as shown above. The 8-step, For FL2VA, change `task` and add `input_reference` as shown above. The 8-step,
ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is
dynamic-only and does not support prefusion, DLO, or LoRA composition. dynamic-only and does not support prefusion or LoRA composition. DLO is
It also rejects model-level, layerwise, and distributed layerwise offload; supported by keeping the request-switchable LoRA A/B buffers resident on the
the legacy dynamic LoRA tensors do not participate in those weight lifecycles. accelerator while DLO streams only the base blocks; budget for this additional
fixed HBM usage. Model-level and standard layerwise offload remain unsupported.
The five requested sigma points produce the four denoiser evaluations expected The five requested sigma points produce the four denoiser evaluations expected
by the Turbo artifact. by the Turbo artifact.
+70 -7
View File
@@ -1,7 +1,7 @@
# NemotronLabs VoiceChat 11B # NemotronLabs VoiceChat 11B
> Offline speech-to-speech: a 16 kHz user utterance in, the agent's spoken reply > Speech-to-speech on a 3-stage vLLM-Omni pipeline: offline single-turn
> (text + 22.05 kHz WAV) out, on a 3-stage vLLM-Omni pipeline. > inference and experimental model-native, frame-locked Realtime serving.
## Summary ## Summary
@@ -11,8 +11,8 @@
timeline: a Conformer + NemotronH hybrid-Mamba thinker emits one text token timeline: a Conformer + NemotronH hybrid-Mamba thinker emits one text token
per acoustic frame, a Gemma3-1B EAR-TTS talker turns the text timeline into per acoustic frame, a Gemma3-1B EAR-TTS talker turns the text timeline into
31-quantizer RVQ code stacks, and an RVQ-VAE codec decodes them to audio. 31-quantizer RVQ code stacks, and an RVQ-VAE codec decodes them to audio.
- Mode: Offline single-turn inference (batch=1). Online/duplex serving, - Mode: Offline single-turn inference (batch=1), plus experimental native
function calling, and batch>1 are documented follow-ups. duplex Realtime serving at one 80 ms microphone frame per scheduler wake.
- Maintainer: [`@yuekaizhang`](https://github.com/yuekaizhang) - Maintainer: [`@yuekaizhang`](https://github.com/yuekaizhang)
## When to use this recipe ## When to use this recipe
@@ -33,6 +33,11 @@ NeMo modules (`nemo_vendored/`), so no `nemo_toolkit` install is needed.
[`vllm_omni/model_executor/models/nemotron_voicechat/`](../../vllm_omni/model_executor/models/nemotron_voicechat/) [`vllm_omni/model_executor/models/nemotron_voicechat/`](../../vllm_omni/model_executor/models/nemotron_voicechat/)
- Staged pipeline config: - Staged pipeline config:
[`vllm_omni/deploy/nemotron_labs_voicechat.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat.yaml) [`vllm_omni/deploy/nemotron_labs_voicechat.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat.yaml)
- Native duplex config and E2E probe:
[`vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml),
[`tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py`](../../tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py)
- Nightly model-level gate:
[`tests/e2e/online_serving/test_nemotron_voicechat_duplex.py`](../../tests/e2e/online_serving/test_nemotron_voicechat_duplex.py)
- Upstream model card: - Upstream model card:
[`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B) [`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
- Reference implementation: NVIDIA-NeMo/Speech, branch `nemotron-labs-voicechat` - Reference implementation: NVIDIA-NeMo/Speech, branch `nemotron-labs-voicechat`
@@ -60,7 +65,7 @@ output stayed within one word of the reference in testing.
- OS: Linux - OS: Linux
- Python: 3.12 - Python: 3.12
- vLLM version: 0.26.0 - vLLM version: 0.27.0
- vLLM-Omni version or commit: this PR / current `main` - vLLM-Omni version or commit: this PR / current `main`
#### Command #### Command
@@ -88,6 +93,55 @@ The reply text should read as a coherent spoken-style answer to the question in
the input WAV, and the WAV should transcribe to (approximately) the same text the input WAV, and the WAV should transcribe to (approximately) the same text
with any ASR model. with any ASR model.
#### Native duplex serving
The checkpoint does not contain the underlying Nemotron text tokenizer. Set
`NEMOTRON_VOICECHAT_LLM_PATH` to a local snapshot of
`nvidia/NVIDIA-Nemotron-Nano-9B-v2` before starting an offline deployment.
```bash
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1
export NEMOTRON_VOICECHAT_LLM_PATH=/path/to/NVIDIA-Nemotron-Nano-9B-v2
vllm-omni serve /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
--omni \
--served-model-name nemotron-voicechat \
--deploy-config vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml
```
The single-GPU profile runs three eager fp32 engine processes on one device.
In another shell, stream only the user channel of one of NVIDIA's bundled
stereo conversation fixtures:
```bash
python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \
--model nemotron-voicechat \
--input-wav /path/to/NVIDIA-NemotronLabs-VoiceChat-11B/turn_taking.wav \
--input-channel 0 \
--output-dir results/nemotron_voicechat_duplex \
--timeout-s 600
```
The probe requires a completed response, non-silent 22.05 kHz output, and the
advertised native 80 ms append capabilities. Add `--expect-function-call` and
use `tool_call.wav` to validate the function-call channel. To validate the
Realtime tool round trip, also provide the expected arguments and return a
tool result:
```bash
python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \
--model nemotron-voicechat \
--input-wav /path/to/NVIDIA-NemotronLabs-VoiceChat-11B/tool_call.wav \
--output-dir results/nemotron_voicechat_tool_call \
--expect-function-call \
--expected-function-name generate_random_number \
--expected-function-arguments '{"min":1,"max":50}' \
--function-output 20 \
--expected-post-tool-text "random number"
```
#### Notes #### Notes
- Memory usage: the shipped yaml runs all three stages on one GPU - Memory usage: the shipped yaml runs all three stages on one GPU
@@ -109,5 +163,14 @@ with any ASR model.
instead of pausing. instead of pausing.
- The talker's `max_tokens` is 16383 (its stage prompt is one placeholder - The talker's `max_tokens` is 16383 (its stage prompt is one placeholder
token, and the stage context is 16384). token, and the stage context is 16384).
- Known limitations: batch=1 only; offline single turn (no online/duplex - The 80 ms frame period is a model/protocol contract, not a throughput
serving yet); the function-call channel is decoded but not yet acted on. guarantee. In an eager fp32 H200 profile, Stage 2 emitted 80 ms packets at
roughly 0.87--0.88 s intervals (RTF about 10.8--10.9); placing the three
stages on separate GPUs did not remove the serial 8-step EAR-TTS bottleneck.
The current native duplex path is functionally streaming but not wall-clock
realtime.
- Known limitations: batch=1 only. The native duplex deployment allows one
active session and does not support barge-in. Tool execution remains
client-owned: the server emits function-call events, accepts a validated
`function_call_output`, and resumes the live model with the returned result;
it does not execute arbitrary tools itself.
+2 -1
View File
@@ -85,8 +85,9 @@ recipes/
| [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) | | [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
| [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB | | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
| [`Wan-AI/Wan2.1-VACE.md`](./Wan-AI/Wan2.1-VACE.md) | Unified T2V, I2V, V2LF, FLF2V, inpaint, and R2V | 1x RTX 5090 (1.3B) / 1x L40S 48GB with layerwise offload (14B) | | [`Wan-AI/Wan2.1-VACE.md`](./Wan-AI/Wan2.1-VACE.md) | Unified T2V, I2V, V2LF, FLF2V, inpaint, and R2V | 1x RTX 5090 (1.3B) / 1x L40S 48GB with layerwise offload (14B) |
| [`XiaomiMiMo/MiMo-Audio.md`](./XiaomiMiMo/MiMo-Audio.md) | Offline + online omni audio (TTS, ASR, dialogue) | 1x RTX 5090 / 5090D 32GB |
| [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB | | [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB | | [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB / 2x non-standard RTX 4090 48GB (49,140 MiB each) |
| [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB | | [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
| [`JD/JoyAI-VL-Interaction.md`](./JD/JoyAI-VL-Interaction.md) | Real-time streaming video-language interaction (proactive speak/silence/delegate) | 1x GPU 24GB+ | | [`JD/JoyAI-VL-Interaction.md`](./JD/JoyAI-VL-Interaction.md) | Real-time streaming video-language interaction (proactive speak/silence/delegate) | 1x GPU 24GB+ |
+161
View File
@@ -0,0 +1,161 @@
# MiMo-Audio for omni audio generation and understanding
> Offline + online TTS, ASR, audio understanding, and multi-turn dialogue
## Summary
- Vendor: Xiaomi MiMo
- Model: `XiaomiMiMo/MiMo-Audio-7B-Instruct`
- Task: Text-to-speech, speech-to-text, audio understanding, spoken dialogue
- Mode: Offline end-to-end scripts and online OpenAI-compatible chat serving
- Maintainer: Community
## When to use this recipe
Use this recipe when you want a known-good starting point for running
MiMo-Audio on a single **32 GB Blackwell consumer GPU** (RTX 5090 / 5090D).
The bundled deploy profile [`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml)
splits VRAM between the AR thinker (stage 0) and code2wav (stage 1), avoids
FlashAttention-2 on SM 12.x, and keeps the audio tokenizer on CPU to reduce
Stage-1 GPU pressure.
## References
- Upstream model card: <https://huggingface.co/XiaomiMiMo/MiMo-Audio-7B-Instruct>
- Audio tokenizer: <https://huggingface.co/XiaomiMiMo/MiMo-Audio-Tokenizer>
- Offline example:
[`examples/offline_inference/mimo_audio/end2end.py`](../../examples/offline_inference/mimo_audio/end2end.py)
- Online example:
[`examples/online_serving/mimo_audio/README.md`](../../examples/online_serving/mimo_audio/README.md)
- Default deploy config (generic single GPU):
[`vllm_omni/deploy/mimo_audio.yaml`](../../vllm_omni/deploy/mimo_audio.yaml)
- 5090D-tuned deploy config:
[`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml)
- Pipeline / deploy docs:
[Stage configs](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/)
- Related issue or discussion:
[RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
- Tokenizer device alignment fix (required for CPU tokenizer on stage 1):
[PR #6539](https://github.com/vllm-project/vllm-omni/pull/6539)
## Hardware Support
## GPU
### 1x NVIDIA GeForce RTX 5090 / 5090D 32GB
Two-stage async-chunk pipeline (thinker + code2wav) on a single GPU. Tested
offline with `tts_sft`.
#### Environment
- OS: Ubuntu (Linux x86_64)
- Python: 3.12
- PyTorch: 2.11.0+cu129
- Driver / runtime: NVIDIA driver 570.x, CUDA 12.9 runtime (Blackwell SM 12.x)
- GPU: NVIDIA GeForce RTX 5090 D, 32 GB
- vLLM: 0.26.0+cu129
- vLLM-Omni: `main` (tested while [#6539](https://github.com/vllm-project/vllm-omni/pull/6539) was open for CPU-tokenizer support)
Install vLLM-Omni from source and align the vLLM wheel with your CUDA stack.
On Blackwell, prefer the **cu129** vLLM build over cu13 defaults when using
PyTorch 2.11+cu129.
Download or cache:
- `XiaomiMiMo/MiMo-Audio-7B-Instruct`
- `XiaomiMiMo/MiMo-Audio-Tokenizer` (or a local checkout of the tokenizer repo)
Export the runtime knobs below before offline or online runs:
```bash
export MIMO_AUDIO_TOKENIZER_PATH="${MIMO_AUDIO_TOKENIZER_PATH:-XiaomiMiMo/MiMo-Audio-Tokenizer}"
export MIMO_AUDIO_TOKENIZER_DEVICE=cpu
export MIMO_AUDIO_TOKENIZER_CUDA_GRAPH=0
export VLLM_USE_FLASHINFER_SAMPLER=0
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
```
Optional (mirror / download stability):
```bash
export HF_HUB_DISABLE_XET=1
export HF_ENDPOINT=https://hf-mirror.com # if you use a Hugging Face mirror
```
Do **not** set `VLLM_ATTENTION_BACKEND` on vLLM 0.26.x; configure
`attention_backend: TRITON_ATTN` in the deploy YAML instead.
#### Command
Offline TTS validation from the repository root:
```bash
cd examples/offline_inference/mimo_audio
python3 -u end2end.py \
--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
--query-type tts_sft \
--text "The weather is so nice today."
```
Online serving (requires MiMo chat template):
```bash
export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
vllm serve XiaomiMiMo/MiMo-Audio-7B-Instruct --omni \
--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \
--served-model-name "MiMo-Audio-7B-Instruct" \
--port 18091 \
--chat-template ./examples/online_serving/mimo_audio/chat_template.jinja
```
Pass the same environment exports to the server process when using a CPU
tokenizer or disabling FlashInfer sampling.
#### Verification
Offline:
```bash
ls -lh output_audio/tts_sft/*.wav
```
Expect a non-empty `.wav` under `output_audio/tts_sft/` and a companion `.txt`
with the decoded text. A successful run completes both stages without a
Stage-1 `decode_vq` device mismatch.
Online:
```bash
curl -s -o /tmp/mimo_health.txt -w '%{http_code}' http://127.0.0.1:18091/health
```
Then use the client under
[`examples/online_serving/mimo_audio/`](../../examples/online_serving/mimo_audio/).
#### Notes
- Memory usage (observed offline): Stage-0 weights ~16.6 GiB; KV cache ~4.9 GiB
with `gpu_memory_utilization: 0.78` on stage 0 and `0.12` on stage 1.
- Key flags:
- `--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml` for 5090D memory
and `TRITON_ATTN`.
- `MIMO_AUDIO_TOKENIZER_PATH` is mandatory.
- `MIMO_AUDIO_TOKENIZER_DEVICE=cpu` saves VRAM; requires
[#6539](https://github.com/vllm-project/vllm-omni/pull/6539) (or equivalent
fix) so code2wav decode uses the tokenizer device.
- `VLLM_USE_FLASHINFER_SAMPLER=0` avoids FlashInfer top-p/top-k sampler
issues on consumer GPUs.
- Known limitations:
- Blackwell may log `SM 12.x requires CUDA >= 12.9` from some optional
tooling; the tested path uses TRITON attention and cu129 wheels.
- Without a working `flash-attn` build, audio quality may be slightly
metallic; install FlashAttention when a compatible wheel is available.
- First request triggers Triton JIT for rotary / attention kernels (latency
spike); subsequent requests are faster.
- For longer outputs or batching, you may need to raise `max_model_len` in
the deploy YAML and match `max_position_embeddings` in the model config
(see the offline example README).
+135 -1
View File
@@ -1,4 +1,4 @@
# GLM-TTS for Chinese/English TTS on 1x GPU # GLM-TTS for Chinese/English TTS
## Summary ## Summary
@@ -103,3 +103,137 @@ python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
- Key flags: `--omni` is required; `--trust-remote-code` is needed for the GLM-TTS phoneme tokenizer; the DiT stage enables bucketed model-internal CUDA graphs with eager fallback. - Key flags: `--omni` is required; `--trust-remote-code` is needed for the GLM-TTS phoneme tokenizer; the DiT stage enables bucketed model-internal CUDA graphs with eager fallback.
- Voice cloning: requires `ref_audio` + `ref_text` together. Reference audio should be 3-10 seconds. Feature extraction (WhisperVQ tokenizer, CampPlus ONNX, mel) runs on the model side. - Voice cloning: requires `ref_audio` + `ref_text` together. Reference audio should be 3-10 seconds. Feature extraction (WhisperVQ tokenizer, CampPlus ONNX, mel) runs on the model side.
- Known limitations: First request may be slow due to lazy model loading (WhisperVQ, CampPlus ONNX). Warm-cache RTF is approximately 0.6-0.7x on A40. - Known limitations: First request may be slow due to lazy model loading (WhisperVQ, CampPlus ONNX). Warm-cache RTF is approximately 0.6-0.7x on A40.
### 2x non-standard RTX 4090 48GB
This configuration was personally validated on two non-standard RTX 4090 cards
reporting 49,140 MiB each in `nvidia-smi`. Stock 24GB RTX 4090 cards were not
validated.
#### Environment
- OS: Ubuntu 22.04.1 LTS
- NVIDIA driver: 580.76.05
- Python: 3.12.13
- PyTorch: 2.11.0+cu130
- CUDA runtime and JIT toolkit: 13.0
- vLLM: 0.26.0
- vLLM-Omni: commit `66a9b0c8`, version
`0.26.0rc2.dev5+g66a9b0c84`
- FlashInfer: 0.6.14
The host's default `nvcc` was 11.8. The working CUDA toolkit came from the
Python environment, so `CUDA_HOME` was set to
`$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/cu13`.
#### Command
From the repository root, start the async-chunk server with:
```bash
export CUDA_HOME="$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/cu13"
export PATH="$CUDA_HOME/bin:$PATH"
export LD_LIBRARY_PATH="$CUDA_HOME/lib:${LD_LIBRARY_PATH:-}"
CUDA_VISIBLE_DEVICES=0,1 vllm-omni serve zai-org/GLM-TTS \
--omni \
--trust-remote-code \
--port 8091 \
--deploy-config vllm_omni/deploy/glm_tts.yaml \
--stage-overrides '{"1":{"devices":"1"}}' \
--allowed-local-media-path "${PWD}/tests/assets"
```
Stage 0 runs on GPU 0, while the override places stage 1 on GPU 1. For the sync
variant, add `--no-async-chunk` to the same command.
#### Verification
Run the Chinese voice-cloning check:
```bash
python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
--model zai-org/GLM-TTS \
--text "今天天气很好,我们一起学习语音模型推理。" \
--ref-audio "file://${PWD}/tests/assets/glm_tts/jiayan_zh.wav" \
--ref-text "他当时还跟线下其他的站姐吵架,然后,打架进局子了。" \
--response-format wav \
--output glm_tts_zh.wav
```
Run the English voice-cloning check with the same Chinese reference:
```bash
python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
--model zai-org/GLM-TTS \
--text "Hello, this is a voice cloning test with GLM-TTS." \
--ref-audio "file://${PWD}/tests/assets/glm_tts/jiayan_zh.wav" \
--ref-text "他当时还跟线下其他的站姐吵架,然后,打架进局子了。" \
--response-format wav \
--output glm_tts_en.wav
```
Inspect either output:
```bash
ffprobe -v error -select_streams a:0 \
-show_entries stream=codec_name,sample_rate,channels \
-of default=noprint_wrappers=1 glm_tts_zh.wav
```
Expected output:
```text
codec_name=pcm_s16le
sample_rate=24000
channels=1
```
The English smoke test produced a valid WAV with an actual duration of 5.020
seconds.
#### Observed warm-cache results
These observations are from one controlled run, not a hardware or performance
guarantee. Cold-start and JIT costs were excluded. Values are mean +/- standard
deviation.
##### Measurement procedure
- For the single-GPU server, use the documented command with
`CUDA_VISIBLE_DEVICES=0` and omit `--stage-overrides`; both bundled stages
remain on GPU 0. For the dual-GPU server, use the command as shown.
- Async chunking is the default. Add `--no-async-chunk` for sync mode.
- For non-streaming measurements, use the exact Chinese client command above.
For streaming, add `--stream` to that command and change the output filename
to `glm_tts_zh.pcm`; streaming output is raw PCM, not WAV.
- The fixed Chinese prompt and reference were exactly those shown in
Verification and yielded 3.900 seconds of audio. Each case used one warm-up
followed by three measured requests at concurrency 1.
- Wall time starts immediately before the HTTP POST and stops after the complete
response body. TTFA runs from immediately before the POST to the first
non-empty audio bytes. RTF is wall time divided by audio duration.
- Peak total device memory used was sampled every 0.2 seconds and reported by
NVML via `nvmlDeviceGetMemoryInfo.used`. The readings include background and
baseline use and are not baseline-subtracted.
| GPUs | Mode | Response | Wall time (s) | TTFA (s) | RTF |
|---:|---|---|---:|---:|---:|
| 1 | Async | Non-streaming | 2.144 +/- 0.076 | -- | 0.550 +/- 0.019 |
| 2 | Async | Non-streaming | 1.969 +/- 0.130 | -- | 0.505 +/- 0.033 |
| 1 | Async | Streaming | 2.135 +/- 0.093 | 1.245 +/- 0.178 | 0.548 +/- 0.024 |
| 2 | Async | Streaming | 1.891 +/- 0.006 | 1.170 +/- 0.042 | 0.485 +/- 0.002 |
| 1 | Sync | Non-streaming | 1.822 +/- 0.016 | -- | 0.467 +/- 0.004 |
| 2 | Sync | Non-streaming | 1.792 +/- 0.023 | -- | 0.459 +/- 0.006 |
Moving stage 1 to GPU 1 reduced async non-streaming wall time by 8.2%. It
reduced async streaming wall time by 11.4% and TTFA by 6.0%, and reduced sync
wall time by 1.6%. Single- and dual-GPU outputs were bitwise identical only
within matching execution and response modes.
| Configuration | GPU 0 | GPU 1 |
|---|---:|---:|
| Single-GPU async | 31.52 GiB | -- |
| Dual-GPU async | 29.23 GiB | 2.91 GiB |
| Single-GPU sync | 31.47 GiB | -- |
| Dual-GPU sync | 29.23 GiB | 2.91 GiB |
+2
View File
@@ -13,6 +13,8 @@ transformers >= 5.10.1, < 5.15
av>=14.0.0 av>=14.0.0
omegaconf>=2.3.0 omegaconf>=2.3.0
diffusers==0.40.0 diffusers==0.40.0
# Hub kernels, including LTX-2.5 NATTEN
kernels==0.14.1
safetensors>=0.8.0 safetensors>=0.8.0
accelerate==1.12.0 accelerate==1.12.0
soundfile>=0.13.1 soundfile>=0.13.1
@@ -0,0 +1,143 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import hashlib
import json
import sys
from types import SimpleNamespace
import pytest
from vllm_omni.benchmarks.duplex import omni_duplex_eval_eval as eval_module
from vllm_omni.benchmarks.duplex.omni_duplex_eval_clock import (
extract_timed_sentences,
normalize_response_items,
split_text,
validate_clock,
)
from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import (
DuplexSample,
canonical_task_type,
family_for_split,
load_samples,
task_type_for_split,
)
from vllm_omni.benchmarks.duplex.omni_duplex_eval_metrics import (
PROTOCOL_PIN,
build_content_prompt,
build_reminder_prompt,
build_temporal_prompt,
parse_judge_json,
reminder_window,
summarize_pr_results,
temporal_window,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_dataset_mapping_and_manifest(tmp_path):
manifest = tmp_path / "samples.json"
manifest.write_text(json.dumps([{"id": "a", "split": "RTD_OCR", "video": "clip.mp4"}]), encoding="utf-8")
sample = load_samples(manifest, media_root=tmp_path)[0]
assert sample.family == "rtd"
assert sample.video == str(tmp_path / "clip.mp4")
assert family_for_split("PR_event_reminder") == "pr"
assert canonical_task_type("post-event-reminder") == "post_event_reminder"
assert task_type_for_split("PR_correction") == "correction"
assert load_samples([{"id": "pr", "split": "PR_correction"}])[0].task_type == "correction"
def test_hugging_face_subset_names_are_splits(monkeypatch):
calls = []
def fake_load_dataset(name, **kwargs):
calls.append((name, kwargs))
return [{"id": "pr", "question_text": "Correct this."}]
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
sample = load_samples("Hothan/Omni-DuplexEval", split="PR_correction")[0]
assert calls == [("Hothan/Omni-DuplexEval", {"split": "PR_correction"})]
assert sample.split == "PR_correction"
assert sample.task_type == "correction"
def test_response_aliases_and_clock_guard():
assert split_text("One. Two!") == ["One.", "Two!"]
assert normalize_response_items({"chunks": [{"text": "x", "current_time": 800}]}) == [
{"sentence": "x", "start": 800.0, "end": 800.0}
]
with pytest.raises(ValueError, match="clock=invalid"):
validate_clock({"clock": "invalid"})
validate_clock({"clock": "invalid"}, allow_invalid=True)
timed = extract_timed_sentences([{"type": "response.output_text.delta", "delta": "Done.", "_media_clock_ms": 800}])
assert [item.as_dict() for item in timed] == [{"sentence": "Done.", "start": 0.8, "end": 0.8}]
with pytest.raises(ValueError, match="only clock=media"):
extract_timed_sentences([], clock="wall")
def test_protocol_windows_and_parsing():
assert PROTOCOL_PIN == "ca3c122b4d4bf67afd6b18ea5e724b4561bdde48"
assert temporal_window(4, 5, 10) == (2.0, 3.0)
assert temporal_window(1, 1.2, 10) == (0.0, 0.5)
assert reminder_window(5) == (5.0, 15.0)
assert parse_judge_json('noise {"content_score": 2.345, "is_relevant": 1}') == {
"content_score": 2.35,
"is_relevant": 1,
}
assert parse_judge_json('noise "success_score": 1') == {"success_score": 1}
assert summarize_pr_results([{"task_type": "correction", "all_success": 1}])["mean_all_success"] == 1.0
def test_protocol_prompt_wording_is_pinned():
prompts = [
build_temporal_prompt(1, 2, "response", "question"),
build_content_prompt("response", "question", ["one", "two"]),
build_reminder_prompt("instruction", "response", "proactive_reminder"),
build_reminder_prompt("instruction", "response", "post_event_reminder"),
build_reminder_prompt("instruction", "response", "correction", "answer"),
]
assert [hashlib.sha256(prompt.encode()).hexdigest() for prompt in prompts] == [
"5f7f5fcb4a37e4d2eea8ba23beabf74630e6fc2fb5b2661ceb94403254bb04d5",
"1fccbed6c1ed8ea1c76d5564d7289ed6adca3678f6a0da2cc735942123d64457",
"d0a0e3a9d873d2f368c5c99b9fe46681f683cc36a9fe0e09b99d6cbca1d6dd0c",
"08f9951243c859d271f82ee53da22cd0bdab1fe93c6ca04f4870d3b959c0803a",
"3de131a7e36cb4b8ac369bb05e87ede518be20221e214df3644211c179131a73",
]
def test_frame_sample_content_passes_frames(tmp_path, monkeypatch):
response = tmp_path / "response.json"
response.write_text(json.dumps([{"sentence": "A person moves.", "start": 1, "end": 2}]), encoding="utf-8")
response.with_name("response.meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8")
video = tmp_path / "video.mp4"
video.write_bytes(b"video")
monkeypatch.setattr(eval_module, "extract_jpeg", lambda *args, **kwargs: b"\xff\xd8jpeg")
class Judge:
def temporal(self, *args, **kwargs):
return '{"temporal_score": 3, "is_relevant": 1}'
def content(self, prompt, video, frames=None, *, mode="video_url"):
assert mode == "frame-sample"
assert frames == [b"\xff\xd8jpeg", b"\xff\xd8jpeg"]
return '{"content_score": 3}'
sample = DuplexSample(
id="rtd",
split="RTD_OCR",
family="rtd",
task_type=None,
video=video,
video_duration=4,
)
score = eval_module.evaluate_sample(
sample,
response,
tmp_path / "score.json",
Judge(),
judge_video_mode="frame-sample",
)
assert score["content"]["frame_count"] == 2
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for Daily-Omni ``--daily-omni-pack-mode`` packing recipes. """Tests for Daily-Omni ``--daily-omni-pack-mode`` packing recipes.
The ``minicpm-interleave`` mode is an accuracy-critical protocol: MiniCPM-o only reaches The ``minicpm-interleave`` mode is an accuracy-critical protocol: MiniCPM-o only reaches
@@ -27,6 +30,8 @@ _MODULE_NAME = "vllm_omni.benchmarks.data_modules.daily_omni_dataset"
if _MODULE_NAME not in sys.modules: if _MODULE_NAME not in sys.modules:
_spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH) _spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
if _spec is None or _spec.loader is None:
raise ImportError(f"Cannot load {_MODULE_NAME} from {_MODULE_PATH}")
_mod = importlib.util.module_from_spec(_spec) _mod = importlib.util.module_from_spec(_spec)
sys.modules[_MODULE_NAME] = _mod sys.modules[_MODULE_NAME] = _mod
_spec.loader.exec_module(_mod) _spec.loader.exec_module(_mod)
@@ -87,7 +92,9 @@ def _stub_extract(monkeypatch, dataset: DailyOmniDataset, counter: list[int]) ->
def fake_extract(video_path, *, audio_path, include_audio, **_kw): def fake_extract(video_path, *, audio_path, include_audio, **_kw):
counter.append(1) counter.append(1)
frames = [Image.new("RGB", (16, 16), color=(i, i, i)) for i in range(_NUM_FRAMES)] frames = [Image.new("RGB", (16, 16), color=(i, i, i)) for i in range(_NUM_FRAMES)]
segments = [np.zeros(1600, dtype=np.float32) for _ in range(_NUM_FRAMES)] if include_audio else [] segments: list[np.ndarray] = (
[np.zeros(1600, dtype=np.float32) for _ in range(_NUM_FRAMES)] if include_audio else []
)
return frames, segments return frames, segments
monkeypatch.setattr(dataset, "_extract_minicpm_frame_audio_segments", fake_extract) monkeypatch.setattr(dataset, "_extract_minicpm_frame_audio_segments", fake_extract)
+190 -16
View File
@@ -34,6 +34,7 @@ from vllm_omni.config.stage_config import (
build_stage_runtime_overrides, build_stage_runtime_overrides,
load_deploy_config, load_deploy_config,
merge_pipeline_deploy, merge_pipeline_deploy,
normalize_pipeline_cli_overrides,
pipeline_cfg_resolver, pipeline_cfg_resolver,
) )
from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
@@ -429,6 +430,20 @@ class TestStageConfigFactory:
assert cfg["final_output"] is True assert cfg["final_output"] is True
assert cfg["final_output_type"] == "image" assert cfg["final_output_type"] == "image"
@pytest.mark.parametrize(
"model_class_name",
[
"MiniMaxH3Pipeline",
"WanPipeline",
"LTX2Pipeline",
"Cosmos3OmniDiffusersPipeline",
],
)
def test_default_video_diffusion_output_type(self, model_class_name):
configs = StageConfigFactory.create_default_diffusion({"model_class_name": model_class_name})
assert configs[0]["final_output_type"] == "video"
def test_default_diffusion_with_parallel_config(self): def test_default_diffusion_with_parallel_config(self):
"""Test diffusion config calculates devices from parallel_config.""" """Test diffusion config calculates devices from parallel_config."""
@@ -534,6 +549,7 @@ class TestPipelineDiscovery:
"""Check that specific models are in OMNI_PIPELINES.""" """Check that specific models are in OMNI_PIPELINES."""
assert "qwen2_5_omni" in OMNI_PIPELINES assert "qwen2_5_omni" in OMNI_PIPELINES
assert "qwen3_omni_moe" in OMNI_PIPELINES assert "qwen3_omni_moe" in OMNI_PIPELINES
assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES
assert "qwen3_tts" in OMNI_PIPELINES assert "qwen3_tts" in OMNI_PIPELINES
def test_registry_resolver_qwen3_omni_all_stages(self): def test_registry_resolver_qwen3_omni_all_stages(self):
@@ -558,6 +574,18 @@ class TestPipelineDiscovery:
assert pipeline.default_deploy_config_name is None assert pipeline.default_deploy_config_name is None
assert len(pipeline.stages) == 1 # thinker only assert len(pipeline.stages) == 1 # thinker only
def test_registry_qwen3_omni_thinker_only_key_is_static(self):
"""Explicit registry key must not re-enter the HF-config resolver."""
pipeline = resolve_pipeline_config(
"qwen3_omni_moe_thinker_only",
Q3_OMNI_ALL_STAGES_HF_CONFIG,
)
assert pipeline is OMNI_PIPELINES["qwen3_omni_moe_thinker_only"]
assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
assert len(pipeline.stages) == 1
assert pipeline.stages[0].engine_output_type == "text"
assert pipeline.stages[0].final_output_type == "text"
@pytest.mark.parametrize( @pytest.mark.parametrize(
"pipeline", "pipeline",
[registered for registered in OMNI_PIPELINES.values() if isinstance(registered, PipelineConfig)], [registered for registered in OMNI_PIPELINES.values() if isinstance(registered, PipelineConfig)],
@@ -1109,8 +1137,35 @@ class TestDeployConfigLoading:
with pytest.raises(ValueError, match=r"stage_args.*PipelineConfig.*stages"): with pytest.raises(ValueError, match=r"stage_args.*PipelineConfig.*stages"):
load_deploy_config(deploy_path) load_deploy_config(deploy_path)
def test_load_minicpmo_duplex_deploy_config(self): @pytest.mark.parametrize(
deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml")) ("filename", "max_sessions"),
[
("minicpmo_4_5.yaml", 4),
("minicpmo_4_5_2gpu.yaml", 4),
("minicpmo_4_5_3gpu.yaml", 4),
("minicpmo_4_5_8x4090.yaml", 1),
("minicpmo_4_5_3gpu_stage1_replicas.yaml", 4),
("minicpmo_4_5_4gpu_stage1_replicas.yaml", 4),
("minicpmo_4_5_8x4090_stage1_replicas.yaml", 1),
],
)
def test_minicpmo_deploy_configs_enable_duplex(self, filename: str, max_sessions: int):
"""Every MiniCPM-o profile serves /v1/realtime, so the retired
minicpmo_4_5_duplex.yaml overlay has no replacement to fall back on.
"""
deploy = load_deploy_config(Path(get_deploy_config_path(filename)))
pipeline = resolve_pipeline_config("minicpmo_4_5")
assert isinstance(pipeline, PipelineConfig)
stages = merge_pipeline_deploy(pipeline, deploy)
assert deploy.session_mode == "duplex"
assert deploy.active_stream_window == 1
assert deploy.duplex_session.max_sessions == max_sessions
assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
def test_load_minicpmo_default_deploy_config(self):
deploy_path = Path(get_deploy_config_path("minicpmo_4_5.yaml"))
deploy = load_deploy_config(deploy_path) deploy = load_deploy_config(deploy_path)
pipeline = resolve_pipeline_config("minicpmo_4_5") pipeline = resolve_pipeline_config("minicpmo_4_5")
@@ -1118,17 +1173,15 @@ class TestDeployConfigLoading:
stages = merge_pipeline_deploy(pipeline, deploy) stages = merge_pipeline_deploy(pipeline, deploy)
assert deploy.session_mode == "duplex"
assert deploy.async_chunk is True assert deploy.async_chunk is True
assert deploy.active_stream_window == 1 assert stages[0].yaml_engine_args["async_scheduling"] is True
assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"] assert stages[1].yaml_engine_args["async_scheduling"] is True
assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"] assert "Async" in (stages[0].scheduler_cls or "")
assert [stage.yaml_engine_args["async_scheduling"] for stage in stages] == [False, False, False] assert "Async" in (stages[1].scheduler_cls or "")
assert all("Async" not in (stage.scheduler_cls or "") for stage in stages)
assert [stage.devices for stage in deploy.stages] == ["0", "0", "0"] assert [stage.devices for stage in deploy.stages] == ["0", "0", "0"]
assert deploy.stages[1].enforce_eager is False assert deploy.stages[1].enforce_eager is False
assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 4096 assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 4096
assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 0 assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 50
assert stages[1].yaml_extras["default_sampling_params"]["temperature"] == 0.8 assert stages[1].yaml_extras["default_sampling_params"]["temperature"] == 0.8
assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [6561] assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [6561]
assert "codec_sampling_params" not in stages[1].yaml_engine_args assert "codec_sampling_params" not in stages[1].yaml_engine_args
@@ -1200,14 +1253,22 @@ stages:
assert deploy.connectors is not None assert deploy.connectors is not None
assert deploy.platforms is not None assert deploy.platforms is not None
def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self): @pytest.mark.parametrize(
("hf_config", "model"),
[
(Q3_OMNI_THINKER_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Thinking"),
(Q3_OMNI_ALL_STAGES_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Instruct"),
],
ids=["thinker_checkpoint", "instruct_forced_thinker_only"],
)
def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self, hf_config, model):
deploy_path = Path(get_deploy_config_path("qwen3_omni_moe_thinking.yaml")) deploy_path = Path(get_deploy_config_path("qwen3_omni_moe_thinking.yaml"))
with patch( with patch(
"vllm_omni.config.config_factory.get_config", "vllm_omni.config.config_factory.get_config",
return_value=Q3_OMNI_THINKER_HF_CONFIG, return_value=hf_config,
): ):
pipeline = StageConfigFactory.get_pipeline_config( pipeline = StageConfigFactory.get_pipeline_config(
"Qwen/Qwen3-Omni-30B-A3B-Thinking", model,
trust_remote_code=False, trust_remote_code=False,
deploy_config_path=str(deploy_path), deploy_config_path=str(deploy_path),
) )
@@ -1216,6 +1277,7 @@ stages:
assert pipeline.model_type == "qwen3_omni_moe_thinker_only" assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
deploy = load_deploy_config(deploy_path) deploy = load_deploy_config(deploy_path)
assert deploy.pipeline == "qwen3_omni_moe_thinker_only"
stages = merge_pipeline_deploy(pipeline, deploy) stages = merge_pipeline_deploy(pipeline, deploy)
assert len(stages) == 1 assert len(stages) == 1
@@ -1267,6 +1329,8 @@ stages:
@pytest.mark.parametrize( @pytest.mark.parametrize(
("deploy_name", "pipeline_name", "stage_count", "final_output_type", "declared_pipeline"), ("deploy_name", "pipeline_name", "stage_count", "final_output_type", "declared_pipeline"),
[ [
("mammoth_moda2.yaml", "mammoth_moda2", 2, "image", "mammoth_moda2"),
("mammoth_moda2_ar.yaml", "mammoth_moda2_ar", 1, "text", "mammoth_moda2_ar"),
("omnivoice.yaml", "omnivoice", 1, "audio", "omnivoice"), ("omnivoice.yaml", "omnivoice", 1, "audio", "omnivoice"),
("mimo_audio.yaml", "mimo_audio", 2, "audio", None), ("mimo_audio.yaml", "mimo_audio", 2, "audio", None),
("step_audio_2.yaml", "step_audio_2", 2, "audio", None), ("step_audio_2.yaml", "step_audio_2", 2, "audio", None),
@@ -1275,6 +1339,13 @@ stages:
("step_audio2_ci.yaml", "step_audio_2_asr", 1, "text", "step_audio_2_asr"), ("step_audio2_ci.yaml", "step_audio_2_asr", 1, "text", "step_audio_2_asr"),
("hunyuan_video_15.yaml", "hunyuan_video_15", 1, "video", None), ("hunyuan_video_15.yaml", "hunyuan_video_15", 1, "video", None),
("wan2_2_ti2v.yaml", "wan2_2_ti2v", 1, "video", None), ("wan2_2_ti2v.yaml", "wan2_2_ti2v", 1, "video", None),
(
"minimax_h3_disaggregated.yaml",
"minimax_h3_disaggregated",
2,
"video",
"minimax_h3_disaggregated",
),
], ],
) )
def test_load_new_registry_backed_deploy_configs( def test_load_new_registry_backed_deploy_configs(
@@ -1295,9 +1366,91 @@ stages:
assert len(stages) == stage_count assert len(stages) == stage_count
assert stages[-1].final_output is True assert stages[-1].final_output is True
assert stages[-1].final_output_type == final_output_type assert stages[-1].final_output_type == final_output_type
if pipeline_name == "minimax_h3_disaggregated":
assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root")
assert (
stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path")
)
assert "model_path_resolver" not in stages[0].yaml_extras
assert "model_path_resolver" not in stages[1].yaml_extras
if deploy.trust_remote_code is not None: if deploy.trust_remote_code is not None:
assert {s.yaml_engine_args.get("trust_remote_code") for s in stages} == {deploy.trust_remote_code} assert {s.yaml_engine_args.get("trust_remote_code") for s in stages} == {deploy.trust_remote_code}
def test_minimax_h3_resolvers_reach_production_engine_args(self, tmp_path: Path, monkeypatch):
from vllm_omni.engine import stage_init_utils
model_root = tmp_path / "MiniMax-H3"
(model_root / "FL2VA" / "text_encoder").mkdir(parents=True)
(model_root / "Ref2VA" / "text_encoder").mkdir(parents=True)
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None)
deploy = load_deploy_config(get_deploy_config_path("minimax_h3_disaggregated.yaml"))
assert deploy.stages[1].engine_extras["model_loaded"] == {"text_encoder": False}
stages = merge_pipeline_deploy(OMNI_PIPELINES["minimax_h3_disaggregated"], deploy)
resolved = [stage_init_utils.build_engine_args_dict(stage.to_omegaconf(), str(model_root)) for stage in stages]
assert resolved[0]["model"] == str(model_root / "FL2VA" / "text_encoder")
assert resolved[1]["model"] == str(model_root / "FL2VA")
assert all("model_path_resolver" not in engine_args for engine_args in resolved)
def test_minimax_h3_text_encoder_tp_reaches_only_stage_zero(self):
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
stages, _ = StageConfigFactory._create_legacy_from_registry(
pipeline,
{"text_encoder_tp_size": 4},
)
assert stages[0].runtime_overrides["tensor_parallel_size"] == 4
assert "text_encoder_tp_size" not in stages[0].runtime_overrides
assert "text_encoder_tp_size" not in stages[1].runtime_overrides
assert stages[1].yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1
def test_minimax_h3_disaggregated_defaults_match_validated_topology(self):
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
deploy = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated.yaml")))
stages = merge_pipeline_deploy(pipeline, deploy)
assert stages[0].yaml_engine_args["max_num_seqs"] == 1
assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root")
assert stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path")
parallel = stages[1].yaml_engine_args["parallel_config"]
assert parallel["tensor_parallel_size"] == 1
assert parallel["ulysses_degree"] == 4
assert parallel["vae_patch_parallel_size"] == 4
assert stages[1].yaml_engine_args["inline_diffusion"] is True
assert stages[1].yaml_extras["default_sampling_params"]["num_inference_steps"] == 50
turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml")))
turbo_stages = merge_pipeline_deploy(pipeline, turbo)
turbo_sampling = turbo_stages[1].yaml_extras["default_sampling_params"]
assert turbo_sampling["num_inference_steps"] == 5
assert turbo_sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0}
def test_minimax_h3_text_encoder_tp_alias_targets_stage_zero(self):
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
assert normalize_pipeline_cli_overrides(pipeline, {}) == {}
assert normalize_pipeline_cli_overrides(pipeline, {"text_encoder_tp_size": 4}) == {
"stage_0_tensor_parallel_size": 4
}
def test_minimax_h3_stage_zero_tp_override_wins_alias_conflict(self):
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
with pytest.warns(UserWarning, match="stage_0_tensor_parallel_size=2 takes precedence"):
normalized = normalize_pipeline_cli_overrides(
pipeline,
{"text_encoder_tp_size": 4, "stage_0_tensor_parallel_size": 2},
)
assert normalized == {"stage_0_tensor_parallel_size": 2}
def test_minimax_h3_rejects_text_encoder_tp_on_diffusion_stage(self):
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
with pytest.raises(ValueError, match="stage_1_text_encoder_tp_size cannot be set"):
normalize_pipeline_cli_overrides(pipeline, {"stage_1_text_encoder_tp_size": 4})
@pytest.mark.parametrize( @pytest.mark.parametrize(
("config_json", "model_index", "expected_pipeline"), ("config_json", "model_index", "expected_pipeline"),
[ [
@@ -1333,6 +1486,29 @@ stages:
assert pipeline is not None assert pipeline is not None
assert pipeline.model_type == expected_pipeline assert pipeline.model_type == expected_pipeline
def test_minimax_h3_disaggregation_is_not_auto_discovered(self):
def get_model_file(filename, _model, revision=None):
del revision
if filename == "config.json":
return None
if filename == "model_index.json":
return {"_class_name": "MiniMaxH3Pipeline"}
return None
with (
patch.object(StageConfigFactory, "get_hf_config", return_value=None),
patch(
"vllm_omni.config.config_factory.get_hf_file_to_dict",
side_effect=get_model_file,
),
):
pipeline = StageConfigFactory.get_pipeline_config(
model="/models/unrelated-checkpoint-name",
trust_remote_code=False,
)
assert pipeline is None
@pytest.mark.parametrize("deploy_name", ["step_audio_2.yaml", "step_audio_2_async_chunk.yaml"]) @pytest.mark.parametrize("deploy_name", ["step_audio_2.yaml", "step_audio_2_async_chunk.yaml"])
def test_step_audio2_deploy_configs_fit_two_gpus(self, deploy_name): def test_step_audio2_deploy_configs_fit_two_gpus(self, deploy_name):
deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name))) deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
@@ -2116,7 +2292,8 @@ class TestBaseConfigInheritance:
def test_minicpmo_overlays_inherit_talker_sampling_params(self): def test_minicpmo_overlays_inherit_talker_sampling_params(self):
"""Overlays must keep the base Talker codec Sampler knobs.""" """Overlays must keep the base Talker codec Sampler knobs."""
for filename in ( for filename in (
"minicpmo_4_5_duplex.yaml", "minicpmo_4_5_2gpu.yaml",
"minicpmo_4_5_3gpu.yaml",
"minicpmo_4_5_3gpu_stage1_replicas.yaml", "minicpmo_4_5_3gpu_stage1_replicas.yaml",
"minicpmo_4_5_4gpu_stage1_replicas.yaml", "minicpmo_4_5_4gpu_stage1_replicas.yaml",
"minicpmo_4_5_8x4090_stage1_replicas.yaml", "minicpmo_4_5_8x4090_stage1_replicas.yaml",
@@ -2128,9 +2305,6 @@ class TestBaseConfigInheritance:
sampling = deploy.stages[1].default_sampling_params sampling = deploy.stages[1].default_sampling_params
assert sampling is not None, f"{filename} stage 1 lost default_sampling_params" assert sampling is not None, f"{filename} stage 1 lost default_sampling_params"
assert sampling["temperature"] == 0.8, filename assert sampling["temperature"] == 0.8, filename
# Upstream utils.TTSSamplingParams. min_tokens is not checked here:
# the duplex overlay drops it to 0 because the model budgets each
# generate_chunk itself.
assert sampling["top_k"] == 25, filename assert sampling["top_k"] == 25, filename
assert sampling["top_p"] == 0.85, filename assert sampling["top_p"] == 0.85, filename
assert sampling["repetition_penalty"] == 1.05, filename assert sampling["repetition_penalty"] == 1.05, filename
+2
View File
@@ -66,6 +66,8 @@ def test_npu_platform_defers_minimax_h3_encoder_patch() -> None:
runtime_src = src[runtime_start:] runtime_src = src[runtime_start:]
assert "apply_minimax_h3_qwen3vl_patch" not in init_src assert "apply_minimax_h3_qwen3vl_patch" not in init_src
assert "apply_minimax_h3_qwen3vl_patch" in runtime_src assert "apply_minimax_h3_qwen3vl_patch" in runtime_src
assert "apply_minimax_h3_qwen3vl_sdpa_patch" not in init_src
assert "apply_minimax_h3_qwen3vl_sdpa_patch" in runtime_src
assert "apply_minimax_h3_qwen3vl_swiglu_patch" in runtime_src assert "apply_minimax_h3_qwen3vl_swiglu_patch" in runtime_src
+3 -3
View File
@@ -164,10 +164,10 @@ def test_inventory_matches_reviewed_snapshot_counts():
"""Make an inventory expansion an explicit review decision.""" """Make an inventory expansion an explicit review decision."""
category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values()) category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values())
assert category_counts == { assert category_counts == {
EnvironmentVariableCategory.PUBLIC_OMNI: 22, EnvironmentVariableCategory.PUBLIC_OMNI: 24,
EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.INHERITED_VLLM: 20,
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
EnvironmentVariableCategory.MODEL_SPECIFIC: 57, EnvironmentVariableCategory.MODEL_SPECIFIC: 56,
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
EnvironmentVariableCategory.INTERNAL: 2, EnvironmentVariableCategory.INTERNAL: 2,
} }
@@ -182,7 +182,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
ModelEnvironmentVariableDisposition.EXTERNAL: 0, ModelEnvironmentVariableDisposition.EXTERNAL: 0,
ModelEnvironmentVariableDisposition.INTERNALIZE: 11, ModelEnvironmentVariableDisposition.INTERNALIZE: 11,
ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 7, ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 6,
} }
+46 -1
View File
@@ -92,10 +92,23 @@ def _from_pipeline_key(
) )
def test_minimax_h3_text_encoder_tp_targets_only_structured_stage_zero() -> None:
config = _from_pipeline_key(
"minimax_h3_disaggregated",
cli_overrides={"text_encoder_tp_size": 4},
)
assert config.stage_by_id(0).parallel_config.tensor_parallel_size == 4
diffusion = config.stage_by_id(1)
assert diffusion.parallel_config.tensor_parallel_size == 1
assert diffusion.parallel_config.text_encoder_tp_size == 1
def test_duplex_session_capacity_propagates_to_every_structured_stage() -> None: def test_duplex_session_capacity_propagates_to_every_structured_stage() -> None:
omni_config = _from_pipeline_key("personaplex") omni_config = _from_pipeline_key("personaplex")
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [2, 2] assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [2, 2]
assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["duplex", "duplex"]
def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -> None: def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -> None:
@@ -105,6 +118,7 @@ def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -
omni_config = _from_pipeline_key("personaplex", deploy_config_path=str(deploy_path)) omni_config = _from_pipeline_key("personaplex", deploy_config_path=str(deploy_path))
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [1, 1] assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [1, 1]
assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["turn", "turn"]
@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES)) @pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
@@ -132,6 +146,7 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod
engine_args = legacy_stage.yaml_engine_args engine_args = legacy_stage.yaml_engine_args
assert omni_stage.model_config.duplex_max_sessions == engine_args.get("duplex_max_sessions", 1) assert omni_stage.model_config.duplex_max_sessions == engine_args.get("duplex_max_sessions", 1)
assert omni_stage.model_config.session_mode == engine_args.get("session_mode", "turn")
assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False) assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False)
assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto") assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto")
assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto") assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
@@ -579,6 +594,7 @@ def test_sub_config_fields_match_structured_scopes():
"interleave_mm_strings", "interleave_mm_strings",
"media_io_kwargs", "media_io_kwargs",
"active_stream_window", "active_stream_window",
"session_mode",
"duplex_max_sessions", "duplex_max_sessions",
"enable_sleep_mode", "enable_sleep_mode",
"default_sampling_params", "default_sampling_params",
@@ -639,8 +655,8 @@ def test_sub_config_fields_match_structured_scopes():
"ulysses_mode", "ulysses_mode",
"cfg_parallel_size", "cfg_parallel_size",
"vae_patch_parallel_size", "vae_patch_parallel_size",
"text_encoder_tp_size",
"vae_parallel_mode", "vae_parallel_mode",
"text_encoder_tp_size",
"use_hsdp", "use_hsdp",
"mask_sp_padding", "mask_sp_padding",
"hsdp_shard_size", "hsdp_shard_size",
@@ -936,6 +952,35 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false") omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
def test_from_pipeline_config_routes_ltx2_conv_vae_extra(tmp_path):
deploy_path = tmp_path / "dreamzero_ltx2_extras.yaml"
deploy_path.write_text(
"\n".join(
[
"pipeline: dreamzero",
"async_chunk: false",
"stages:",
" - stage_id: 0",
" extras:",
" ltx2_use_conv_vae: true",
]
)
)
stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True
def test_stage_override_routes_ltx2_conv_vae_extra():
stage = _from_pipeline_key(
"dreamzero",
cli_overrides={"stage_0_extras": {"ltx2_use_conv_vae": True}},
).stage_by_id(0)
assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True
def test_structured_diffusion_config_rejects_invalid_compile_granularity(): def test_structured_diffusion_config_rejects_invalid_compile_granularity():
with pytest.raises(ValidationError, match="diffusion_compile_granularity"): with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block") omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")
+7
View File
@@ -66,3 +66,10 @@ def test_register_resolver_requires_model_type(custom_resolver, clean_pipeline_r
"""Ensure that registering a custom resolver to OMNI_PIPELINES requires an explicit model_type.""" """Ensure that registering a custom resolver to OMNI_PIPELINES requires an explicit model_type."""
with pytest.raises(ValueError): with pytest.raises(ValueError):
register_pipeline(custom_resolver) register_pipeline(custom_resolver)
def test_minimax_h3_disaggregation_is_explicit_opt_in():
assert "minimax_h3" not in OMNI_PIPELINES
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
assert isinstance(pipeline, PipelineConfig)
assert pipeline.model_type == "minimax_h3_disaggregated"
@@ -100,6 +100,7 @@ def _make_drain_sched(session: Request) -> MagicMock:
sched.pending_stop_after_extraction = set() sched.pending_stop_after_extraction = set()
sched.connector = None sched.connector = None
sched.kv_cache_manager.take_events.return_value = None sched.kv_cache_manager.take_events.return_value = None
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
sched.finished_req_ids_dict = {} sched.finished_req_ids_dict = {}
sched.make_stats.return_value = None sched.make_stats.return_value = None
return sched return sched
@@ -25,10 +25,10 @@ from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
pytestmark = [pytest.mark.core_model, pytest.mark.cpu] pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler: def _make_scheduler(*, stage_id: int = 0, session_mode: str = "turn") -> OmniARScheduler:
sched = OmniARScheduler.__new__(OmniARScheduler) sched = OmniARScheduler.__new__(OmniARScheduler)
sched._new_prompt_len_snapshot = {} sched._new_prompt_len_snapshot = {}
sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id)) sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id, session_mode=session_mode))
sched.num_waiting_for_streaming_input = 0 sched.num_waiting_for_streaming_input = 0
sched.log_stats = False sched.log_stats = False
sched.chunk_transfer_adapter = None sched.chunk_transfer_adapter = None
@@ -93,6 +93,7 @@ def _run_resumable_segment_stop(
sched.pending_stop_after_extraction = set() sched.pending_stop_after_extraction = set()
sched.connector = None sched.connector = None
sched.kv_cache_manager.take_events.return_value = None sched.kv_cache_manager.take_events.return_value = None
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
sched.finished_req_ids_dict = {} sched.finished_req_ids_dict = {}
sched.make_stats.return_value = None sched.make_stats.return_value = None
@@ -198,6 +199,7 @@ def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> Non
inter_stage_output, inter_stage_output,
session, session,
True, True,
new_token_ids=[42],
confirmed_num_computed_tokens=26, confirmed_num_computed_tokens=26,
) )
@@ -226,6 +228,7 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
sched.pending_stop_after_extraction = set() sched.pending_stop_after_extraction = set()
sched.connector = None sched.connector = None
sched.kv_cache_manager.take_events.return_value = None sched.kv_cache_manager.take_events.return_value = None
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
sched.finished_req_ids_dict = {} sched.finished_req_ids_dict = {}
sched.make_stats.return_value = None sched.make_stats.return_value = None
@@ -252,6 +255,37 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
sched.chunk_transfer_adapter.save_async.assert_not_called() sched.chunk_transfer_adapter.save_async.assert_not_called()
def test_queued_streaming_update_on_async_stop_fences_in_flight_once() -> None:
"""Native duplex + async: the next unit is usually already queued when the
current one stops, so _handle_stopped_request applies the update and the
stop site below it fences the same in-flight decode a second time. An
accumulated residue outlives the drain and swallows the next unit's
listen/speak, which the client sees as silence until its timeout.
"""
session = _make_request()
session.status = RequestStatus.RUNNING
session.resumable = True
session.append_output_token_ids([7])
session.num_computed_tokens = 4
session.num_output_placeholders = 1
# This step's frame (settled to 0 by update_from_output) plus one extra
# async decode still in flight.
session.num_in_flight_tokens = 2
sched = _make_scheduler(stage_id=0)
sched._enqueue_waiting_request = MagicMock()
def handle_stopped(request: Request) -> bool:
sched._update_request_as_session(request, _make_update([10, 20]))
return False
_run_resumable_segment_stop(session, handle_stopped=handle_stopped)
# Exactly the one unreported decode, so the drain reaches zero before the
# new segment's first frame arrives.
assert session.num_stale_output_tokens == session.num_in_flight_tokens == 1
def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None: def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
sched = _make_scheduler(stage_id=0) sched = _make_scheduler(stage_id=0)
session = _make_request() session = _make_request()
@@ -396,3 +430,71 @@ def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None:
assert sched.chunk_transfer_adapter.replaced_streaming_prompt_ids == set() assert sched.chunk_transfer_adapter.replaced_streaming_prompt_ids == set()
assert sched.chunk_transfer_adapter.requests_with_ready_chunks == {session.request_id} assert sched.chunk_transfer_adapter.requests_with_ready_chunks == {session.request_id}
assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {} assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {}
def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
"""A WAITING_FOR_CHUNK stop must not delete its newly parked session."""
session = _make_request()
session.resumable = True
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
def queue(*requests):
result = MagicMock()
result.requests = set(requests)
result.add_request.side_effect = result.requests.add
result.remove_requests.side_effect = result.requests.difference_update
return result
sched = OmniARScheduler.__new__(OmniARScheduler)
sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode="duplex"))
sched.running = []
sched.waiting = queue()
sched.skipped_waiting = queue(session)
sched.num_waiting_for_streaming_input = 1
sched.chunk_transfer_adapter = SimpleNamespace(
receives_chunks=True,
segment_finished_requests={session.request_id},
)
sched._resume_downstream_chunk_receiver(session)
sched._remove_stopped_requests_from_queues(set(), {session})
assert session.status == RequestStatus.WAITING
assert session in sched.waiting.requests
assert session not in sched.skipped_waiting.requests
assert sched.num_waiting_for_streaming_input == 0
assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
@pytest.mark.parametrize(
("receives_chunks", "session_mode"),
[
(False, "duplex"),
(True, "turn"),
],
)
def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
receives_chunks: bool,
session_mode: str,
) -> None:
"""Only duplex connector-driven receivers resume without an update."""
session = _make_request()
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
sched = OmniARScheduler.__new__(OmniARScheduler)
sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode=session_mode))
sched.num_waiting_for_streaming_input = 1
sched.chunk_transfer_adapter = SimpleNamespace(
receives_chunks=receives_chunks,
segment_finished_requests={session.request_id},
)
sched.skipped_waiting = MagicMock()
sched._enqueue_waiting_request = MagicMock()
sched._resume_downstream_chunk_receiver(session)
assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
assert sched.num_waiting_for_streaming_input == 1
assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
sched.skipped_waiting.remove_requests.assert_not_called()
sched._enqueue_waiting_request.assert_not_called()
@@ -138,6 +138,7 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
sched.recompute_kv_load_failures = False sched.recompute_kv_load_failures = False
sched.connector = None sched.connector = None
sched.kv_cache_manager.take_events.return_value = None sched.kv_cache_manager.take_events.return_value = None
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
sched.finished_req_ids_dict = {} sched.finished_req_ids_dict = {}
sched.make_stats.return_value = None sched.make_stats.return_value = None
@@ -0,0 +1,168 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""v0.28 deferred KV-block free fence in the Omni schedulers (#6606 review).
Upstream v0.28 (defer_block_free) advances ``sched_step_seq`` for every
non-empty scheduled step and, at the top of ``update_from_output()``,
advances ``processed_step_seq`` and drains the deferred frees. The Omni
schedulers each carried only half of that contract:
- ``OmniARScheduler.schedule()`` delegates to ``super().schedule()`` (fence
advances) but its reimplemented ``update_from_output()`` never advanced
``processed_step_seq`` nor drained -> deferred blocks accumulate without
bound under async/PP KV-consumer configurations.
- ``OmniGenerationScheduler`` reimplements both: the fast-path
``schedule()`` never advanced ``sched_step_seq`` (so finished requests
freed blocks an in-flight GPU step could still be writing) and its
``update_from_output()`` never drained.
These tests pin the restored contract at the call boundary; the drain and
free logic itself is upstream-tested.
"""
from __future__ import annotations
from unittest.mock import MagicMock
import pytest
# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
# before Request / RequestStatus are bound in this module.
# isort: off
import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
from vllm.sampling_params import SamplingParams
from vllm.v1.core.sched.interface import PauseState
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.core.sched.request_queue import SchedulingPolicy
from vllm.v1.outputs import ModelRunnerOutput
from vllm.v1.request import Request, RequestStatus
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
# isort: on
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _make_update_sched(scheduler_cls) -> MagicMock:
"""MagicMock scheduler driving the real ``update_from_output`` with an
empty batch; only the fence state carries real values."""
sched = MagicMock()
sched.requests = {}
sched.perf_metrics = None
sched.chunk_transfer_adapter = None
sched.defer_block_free = True
sched.processed_step_seq = 0
return sched
def _empty_frame(*, total_num_scheduled_tokens: int):
scheduler_output = MagicMock(spec=SchedulerOutput)
scheduler_output.num_scheduled_tokens = {}
scheduler_output.total_num_scheduled_tokens = total_num_scheduled_tokens
model_runner_output = MagicMock(spec=ModelRunnerOutput)
model_runner_output.sampled_token_ids = []
model_runner_output.logprobs = None
model_runner_output.prompt_logprobs_dict = {}
model_runner_output.pooler_output = None
model_runner_output.num_nans_in_logits = None
model_runner_output.kv_connector_output = None
model_runner_output.cudagraph_stats = None
return scheduler_output, model_runner_output
@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
def test_update_from_output_advances_fence_and_drains(scheduler_cls) -> None:
sched = _make_update_sched(scheduler_cls)
scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3)
scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
assert sched.processed_step_seq == 1
sched._drain_deferred_frees.assert_called_once_with()
@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
def test_update_from_output_zero_token_step_leaves_fence(scheduler_cls) -> None:
"""0-token steps do not advance the fence upstream (their seq was never
advanced in schedule()); draining on them would run ahead of the GPU."""
sched = _make_update_sched(scheduler_cls)
scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=0)
scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
assert sched.processed_step_seq == 0
sched._drain_deferred_frees.assert_not_called()
@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
def test_update_from_output_defer_disabled_leaves_fence(scheduler_cls) -> None:
sched = _make_update_sched(scheduler_cls)
sched.defer_block_free = False
scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3)
scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
assert sched.processed_step_seq == 0
sched._drain_deferred_frees.assert_not_called()
def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request]:
"""MagicMock scheduler that drives the real generation fast-path
``schedule()`` over one running one-shot request."""
request = Request(
request_id="req-fence",
prompt_token_ids=[1, 2, 3],
sampling_params=SamplingParams(max_tokens=4),
pooling_params=None,
arrival_time=100.0,
block_hasher=None,
)
request.status = RequestStatus.RUNNING
sched = MagicMock()
sched.max_num_scheduled_tokens = 8
sched._pause_state = PauseState.UNPAUSED
sched.requests = {request.request_id: request}
sched.running = [request]
sched.waiting = [] # falsy: the waiting loop never runs
sched.policy = SchedulingPolicy.FCFS
sched.chunk_transfer_adapter = None
sched.scheduler_config = MagicMock(enable_chunked_prefill=True)
sched.num_lookahead_tokens = 0
sched.log_stats = False
sched.kv_cache_config = MagicMock(kv_cache_groups=[])
sched.use_v2_model_runner = False
sched.prev_step_scheduled_req_ids = set()
sched.needs_kv_cache_zeroing = False
sched.connector = None
sched.ec_connector = None
sched.defer_block_free = defer_block_free
sched.sched_step_seq = 0
# last_sched_seq stamping order: _update_after_schedule must observe the
# already-advanced fence value, as upstream schedule() guarantees.
sched.seq_at_update_after_schedule = None
def _capture(_scheduler_output):
sched.seq_at_update_after_schedule = sched.sched_step_seq
sched._update_after_schedule.side_effect = _capture
return sched, request
def test_fast_path_schedule_advances_fence_before_update_after_schedule() -> None:
sched, request = _make_fast_path_sched(defer_block_free=True)
OmniGenerationScheduler.schedule(sched)
assert sched.sched_step_seq == 1
assert sched.seq_at_update_after_schedule == 1
def test_fast_path_schedule_defer_disabled_leaves_fence() -> None:
sched, request = _make_fast_path_sched(defer_block_free=False)
OmniGenerationScheduler.schedule(sched)
assert sched.sched_step_seq == 0
@@ -0,0 +1,180 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""v0.28 EC request-finish contract in the Omni schedulers (#6606 review).
Upstream v0.28 ``_free_request()`` fires ``ec_connector.request_finished()``
BEFORE the encoder cache is freed (so the connector can inspect per-request
state), honors its delayed-free result, and returns the generated
``ec_transfer_params``; ``update_from_output()`` then carries them into the
``EngineCoreOutput``. ``OmniARScheduler`` previously skipped the hook and
always returned ``None``; both Omni schedulers discarded the second tuple
item, so a remote encoder-cache handle could never reach the frontend.
"""
from __future__ import annotations
from unittest.mock import MagicMock
import pytest
# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
# before Request / RequestStatus are bound in this module.
# isort: off
import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
from vllm.sampling_params import SamplingParams
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.outputs import ModelRunnerOutput
from vllm.v1.request import Request, RequestStatus
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
# isort: on
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
EC_PARAMS = {"remote_encoder_cache_handle": "ec-1"}
def _make_free_request_sched(*, ec_connector) -> tuple[OmniARScheduler, MagicMock]:
"""Minimal real-`_free_request` harness; `order` records the EC hook
firing relative to the encoder-cache free."""
order = MagicMock()
sched = OmniARScheduler.__new__(OmniARScheduler)
sched._omits_kv_transfer_cache = {}
sched._connector_finished = lambda request: (False, None)
sched.ec_connector = ec_connector
if ec_connector is not None:
order.attach_mock(ec_connector.request_finished, "request_finished")
sched.encoder_cache_manager = MagicMock()
order.attach_mock(sched.encoder_cache_manager.free, "encoder_cache_free")
sched.finished_req_ids = set()
sched._new_prompt_len_snapshot = {}
sched.finished_req_ids_dict = None
sched._should_transfer_kv_for_request = lambda req_id: False
sched._free_blocks = MagicMock()
sched._free_input_coordinator_request = MagicMock()
sched.chunk_transfer_adapter = None
return sched, order
class _FakeFinishedRequest:
def __init__(self, request_id: str) -> None:
self.request_id = request_id
def is_finished(self) -> bool:
return True
def test_free_request_fires_ec_hook_before_encoder_cache_free() -> None:
ec_connector = MagicMock()
ec_connector.request_finished.return_value = (False, EC_PARAMS)
sched, order = _make_free_request_sched(ec_connector=ec_connector)
request = _FakeFinishedRequest("req-ec")
kv_params, ec_params = sched._free_request(request)
assert ec_params == EC_PARAMS
ec_connector.request_finished.assert_called_once_with(request)
hook_calls = [name for name, *_ in order.mock_calls]
assert hook_calls.index("request_finished") < hook_calls.index("encoder_cache_free")
sched._free_blocks.assert_called_once_with(request)
def test_free_request_honors_ec_delayed_free() -> None:
ec_connector = MagicMock()
ec_connector.request_finished.return_value = (True, EC_PARAMS)
sched, _ = _make_free_request_sched(ec_connector=ec_connector)
request = _FakeFinishedRequest("req-ec-delay")
kv_params, ec_params = sched._free_request(request)
assert ec_params == EC_PARAMS
sched._free_blocks.assert_not_called()
def test_free_request_without_ec_connector_returns_none() -> None:
sched, _ = _make_free_request_sched(ec_connector=None)
request = _FakeFinishedRequest("req-no-ec")
kv_params, ec_params = sched._free_request(request)
assert ec_params is None
sched._free_blocks.assert_called_once_with(request)
def _make_finishing_request() -> Request:
request = Request(
request_id="req-ec-finish",
prompt_token_ids=[1, 2, 3],
sampling_params=SamplingParams(max_tokens=4),
pooling_params=None,
arrival_time=100.0,
block_hasher=None,
)
request.status = RequestStatus.RUNNING
return request
def _make_finish_sched(request: Request) -> MagicMock:
"""MagicMock scheduler driving the real ``update_from_output`` through a
finishing request; ``_free_request`` reports EC transfer params."""
sched = MagicMock()
sched.requests = {request.request_id: request}
sched.perf_metrics = None
sched.defer_block_free = False
sched.structured_output_manager.should_advance.return_value = False
sched._process_kv_transfer_trigger.return_value = False
sched._maybe_decode_pooling_output.return_value = None
sched._handle_stopped_request.return_value = True
sched._free_request.return_value = (None, EC_PARAMS)
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
sched.chunk_transfer_adapter = None
sched.finished_req_ids_dict = None
sched._new_prompt_len_snapshot = {}
return sched
def test_ar_update_from_output_carries_ec_params_to_engine_output() -> None:
request = _make_finishing_request()
sched = _make_finish_sched(request)
sched._update_request_with_output.return_value = ([42], True)
engine_core_outputs = OmniARScheduler.update_from_output(sched, *_finish_frame_parts(request))
(eco,) = engine_core_outputs[request.client_index].outputs
assert eco.ec_transfer_params == EC_PARAMS
sched._free_request.assert_called_once_with(request)
def _finish_frame_parts(request: Request):
scheduler_output = MagicMock(spec=SchedulerOutput)
scheduler_output.num_scheduled_tokens = {request.request_id: 1}
scheduler_output.total_num_scheduled_tokens = 1
scheduler_output.scheduled_spec_decode_tokens = {}
scheduler_output.num_invalid_spec_tokens = 0
model_runner_output = MagicMock(spec=ModelRunnerOutput)
model_runner_output.sampled_token_ids = [[42]]
model_runner_output.logprobs = None
model_runner_output.prompt_logprobs_dict = {}
model_runner_output.pooler_output = None
model_runner_output.num_nans_in_logits = None
model_runner_output.kv_connector_output = None
model_runner_output.cudagraph_stats = None
model_runner_output.req_id_to_index = {request.request_id: 0}
model_runner_output.routed_experts = None
return scheduler_output, model_runner_output
def test_generation_update_from_output_carries_ec_params_to_engine_output() -> None:
request = _make_finishing_request()
# One-shot generation finish: the whole prompt is computed.
request.num_computed_tokens = len(request.prompt_token_ids)
sched = _make_finish_sched(request)
engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_finish_frame_parts(request))
(eco,) = engine_core_outputs[request.client_index].outputs
assert eco.ec_transfer_params == EC_PARAMS
sched._free_request.assert_called_once_with(request)
@@ -0,0 +1,126 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""v0.28 prefill-stats finalization in the Omni schedulers (#6606 review).
Upstream v0.28 finalizes prefill statistics with
``PrefillStats.finalize(kv_cache_manager.estimate_cached_tokens(request))``
before emitting them, which computes ``num_cache_creation_tokens`` (tokens
computed and written to the prefix cache). Both Omni schedulers previously
emitted ``take_prefill_stats()`` raw, so creation usage was always 0.
"""
from __future__ import annotations
from unittest.mock import MagicMock
import pytest
# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
# before Request / RequestStatus are bound in this module.
# isort: off
import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
from vllm.sampling_params import SamplingParams
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.outputs import ModelRunnerOutput
from vllm.v1.request import Request, RequestStatus
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
# isort: on
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _make_request() -> Request:
request = Request(
request_id="req-prefill-stats",
prompt_token_ids=[1, 2, 3],
sampling_params=SamplingParams(max_tokens=4),
pooling_params=None,
arrival_time=100.0,
block_hasher=None,
)
request.status = RequestStatus.RUNNING
# A scheduled prefill: 3 prompt tokens, 1 from the local prefix cache.
request.prefill_stats.set(
num_prompt_tokens=3,
num_local_cached_tokens=1,
num_external_cached_tokens=0,
)
return request
def _make_sched(request: Request) -> MagicMock:
sched = MagicMock()
sched.requests = {request.request_id: request}
sched.perf_metrics = None
sched.defer_block_free = False
sched.structured_output_manager.should_advance.return_value = False
sched._process_kv_transfer_trigger.return_value = False
sched._maybe_decode_pooling_output.return_value = None
sched._handle_stopped_request.return_value = True
# The whole 3-token prompt ends up cached: 3 - 1 pre-cached = 2 created.
sched.kv_cache_manager.estimate_cached_tokens.return_value = 3
def _free_and_invalidate(req):
# Lifecycle guard: freeing releases the KV blocks, after which the
# manager reports 0 cached tokens. If finalize ran after
# _free_request (the reviewed bug), creation count would compute
# against 0 and the assertions below would see 0, not 2.
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
return (None, None)
sched._free_request.side_effect = _free_and_invalidate
sched.chunk_transfer_adapter = None
sched.finished_req_ids_dict = None
sched._new_prompt_len_snapshot = {}
return sched
def _frame(request: Request):
scheduler_output = MagicMock(spec=SchedulerOutput)
scheduler_output.num_scheduled_tokens = {request.request_id: 1}
scheduler_output.total_num_scheduled_tokens = 1
scheduler_output.scheduled_spec_decode_tokens = {}
scheduler_output.num_invalid_spec_tokens = 0
model_runner_output = MagicMock(spec=ModelRunnerOutput)
model_runner_output.sampled_token_ids = [[42]]
model_runner_output.logprobs = None
model_runner_output.prompt_logprobs_dict = {}
model_runner_output.pooler_output = None
model_runner_output.num_nans_in_logits = None
model_runner_output.kv_connector_output = None
model_runner_output.cudagraph_stats = None
model_runner_output.req_id_to_index = {request.request_id: 0}
model_runner_output.routed_experts = None
return scheduler_output, model_runner_output
def test_ar_scheduler_finalizes_prefill_stats_before_emitting() -> None:
request = _make_request()
sched = _make_sched(request)
sched._update_request_with_output.return_value = ([42], True)
engine_core_outputs = OmniARScheduler.update_from_output(sched, *_frame(request))
(eco,) = engine_core_outputs[request.client_index].outputs
assert eco.prefill_stats is not None
assert eco.prefill_stats.num_cached_tokens == 1
assert eco.prefill_stats.num_cache_creation_tokens == 2
sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request)
def test_generation_scheduler_finalizes_prefill_stats_before_emitting() -> None:
request = _make_request()
request.num_computed_tokens = len(request.prompt_token_ids)
sched = _make_sched(request)
engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_frame(request))
(eco,) = engine_core_outputs[request.client_index].outputs
assert eco.prefill_stats is not None
assert eco.prefill_stats.num_cached_tokens == 1
assert eco.prefill_stats.num_cache_creation_tokens == 2
sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request)
@@ -0,0 +1,106 @@
[
{
"test_name": "test_hunyuanvideo15_i2v_single_device",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100"
},
"num_cards": 1
}
},
"full_model",
"diffusion"
],
"description": "Single-device baseline",
"server_type": "vllm-omni",
"server_params": {
"model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v",
"serve_args": {
"flow-shift": 5.0,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "832x480_frames33_steps4",
"dataset": "random",
"task": "i2v",
"num-prompts": 10,
"max-concurrency": 1,
"num-input-images": 1,
"seed": 42,
"enable-negative-prompt": true,
"random-request-config": [
{
"width": 832,
"height": 480,
"num_inference_steps": 4,
"num_frames": 33,
"fps": 24,
"weight": 1
}
],
"extra-body": {
"guidance_scale": 6.0,
"flow_shift": 5.0
}
}
]
},
{
"test_name": "test_hunyuanvideo15_i2v_cachedit_tp2_vae_patch2",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100"
},
"num_cards": 2
}
},
"full_model",
"diffusion"
],
"description": "CacheDiT + TP=2 + VAE patch parallel=2 + VAE tiling",
"server_type": "vllm-omni",
"server_params": {
"model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v",
"serve_args": {
"flow-shift": 5.0,
"cache-backend": "cache_dit",
"tensor-parallel-size": 2,
"vae-patch-parallel-size": 2,
"vae-use-tiling": true,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "832x480_frames33_steps4",
"dataset": "random",
"task": "i2v",
"num-prompts": 10,
"max-concurrency": 1,
"num-input-images": 1,
"seed": 42,
"enable-negative-prompt": true,
"random-request-config": [
{
"width": 832,
"height": 480,
"num_inference_steps": 4,
"num_frames": 33,
"fps": 24,
"weight": 1
}
],
"extra-body": {
"guidance_scale": 6.0,
"flow_shift": 5.0
}
}
]
}
]
@@ -0,0 +1,104 @@
[
{
"test_name": "test_hunyuanvideo15_t2v_single_device",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100"
},
"num_cards": 1
}
},
"full_model",
"diffusion"
],
"description": "Single-device baseline",
"server_type": "vllm-omni",
"server_params": {
"model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
"serve_args": {
"flow-shift": 5.0,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "832x480_frames33_steps4",
"dataset": "random",
"task": "t2v",
"num-prompts": 10,
"max-concurrency": 1,
"seed": 42,
"enable-negative-prompt": true,
"random-request-config": [
{
"width": 832,
"height": 480,
"num_inference_steps": 4,
"num_frames": 33,
"fps": 24,
"weight": 1
}
],
"extra-body": {
"guidance_scale": 6.0,
"flow_shift": 5.0
}
}
]
},
{
"test_name": "test_hunyuanvideo15_t2v_cachedit_tp2_vae_patch2",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100"
},
"num_cards": 2
}
},
"full_model",
"diffusion"
],
"description": "CacheDiT + TP=2 + VAE patch parallel=2 + VAE tiling",
"server_type": "vllm-omni",
"server_params": {
"model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
"serve_args": {
"flow-shift": 5.0,
"cache-backend": "cache_dit",
"tensor-parallel-size": 2,
"vae-patch-parallel-size": 2,
"vae-use-tiling": true,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "832x480_frames33_steps4",
"dataset": "random",
"task": "t2v",
"num-prompts": 10,
"max-concurrency": 1,
"seed": 42,
"enable-negative-prompt": true,
"random-request-config": [
{
"width": 832,
"height": 480,
"num_inference_steps": 4,
"num_frames": 33,
"fps": 24,
"weight": 1
}
],
"extra-body": {
"guidance_scale": 6.0,
"flow_shift": 5.0
}
}
]
}
]
@@ -18,7 +18,7 @@
"model": "openbmb/MiniCPM-o-4_5", "model": "openbmb/MiniCPM-o-4_5",
"extra_cli_args": [ "extra_cli_args": [
"--deploy-config", "--deploy-config",
"vllm_omni/deploy/minicpmo_4_5_duplex.yaml", "vllm_omni/deploy/minicpmo_4_5.yaml",
"--trust-remote-code" "--trust-remote-code"
] ]
}, },
@@ -308,6 +308,91 @@ def test_flash_attn_func_preferred_over_varlen():
print("✓ flash_attn_func forward works correctly!") print("✓ flash_attn_func forward works correctly!")
@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU")
@pytest.mark.parametrize("k_len", [40, 256])
def test_cross_attn_key_padding_vs_sdpa(k_len):
"""
Case 3: Cross-attention with a key-padding mask.
The mask covers only the key sequence; every query position must attend
to the valid keys. FA and SDPA outputs should be very close. k_len=256
covers equal Q/K lengths, which must not be mistaken for self-attention,
and valid key lengths differ per batch element.
"""
device = torch.device(current_omni_platform.device_type)
dtype = torch.bfloat16
batch_size = 2
q_len = 256
valid_lens = (17, 29)
num_heads = 8
head_dim = 64
fa_impl = FlashAttentionImpl(
num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False, role="cross"
)
sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False)
torch.manual_seed(7)
query = torch.randn(batch_size, q_len, num_heads, head_dim, device=device, dtype=dtype)
key = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype)
value = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype)
attn_mask = torch.zeros(batch_size, k_len, dtype=torch.bool, device=device)
for i, valid_len in enumerate(valid_lens):
attn_mask[i, :valid_len] = True
attn_metadata = AttentionMetadata(attn_mask=attn_mask)
output_fa = fa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata)
output_sdpa = sdpa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata)
max_diff = torch.max(torch.abs(output_fa - output_sdpa)).item()
mean_diff = torch.mean(torch.abs(output_fa - output_sdpa)).item()
print("\n=== Case 3: Cross-Attention Key-Padding FA vs SDPA ===")
print(f"Max absolute difference: {max_diff:.6f}")
print(f"Mean absolute difference: {mean_diff:.6f}")
assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
print("✓ Case 3 PASSED: FA and SDPA cross-attention outputs are very close!")
def test_varlen_masked_routing_by_role(monkeypatch):
"""The unpad route is picked by role, not Q/K length equality; runs without a GPU."""
calls = []
def fake_varlen_func(q, k, v, **kwargs):
calls.append((q.shape[0], kwargs["cu_seqlens_q"], kwargs["cu_seqlens_k"]))
return q
monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func)
batch_size, seq_len, num_heads, head_dim = 2, 4, 2, 8
query = torch.randn(batch_size, seq_len, num_heads, head_dim)
metadata = AttentionMetadata(attn_mask=torch.tensor([[True, True, False, False], [True, True, True, False]]))
cross_impl = FlashAttentionImpl(
num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False, role="cross"
)
output = cross_impl.forward_cuda(query, query, query, metadata)
assert output.shape == query.shape
num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1]
assert num_q_rows == batch_size * seq_len
assert cu_seqlens_q.tolist() == [0, 4, 8]
assert cu_seqlens_k.tolist() == [0, 2, 5]
self_impl = FlashAttentionImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False)
output = self_impl.forward_cuda(query, query, query, metadata)
assert output.shape == query.shape
num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1]
assert num_q_rows == 5
assert cu_seqlens_q.tolist() == [0, 2, 5]
assert cu_seqlens_k.tolist() == [0, 2, 5]
def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch): def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
calls = [] calls = []
+3 -6
View File
@@ -365,13 +365,10 @@ class TestCacheDiTBackend:
@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter") @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit") @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter): def test_enable_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
"""Test DreamID uses pipeline.transformer for cache enable/refresh. """Generic cache path should pick up ``transformer.fused_blocks``."""
NOTE: DreamID no longer has a custom enabler, so this tests against the generic path.
"""
mock_pipeline = Mock() mock_pipeline = Mock()
mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline" mock_pipeline.__class__.__name__ = "FakeFusedBlocksPipeline"
mock_pipeline.transformer = Mock() mock_pipeline.transformer = Mock()
mock_pipeline.transformer.fused_blocks = Mock() mock_pipeline.transformer.fused_blocks = Mock()
mock_pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig( mock_pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
-10
View File
@@ -6,7 +6,6 @@ Model specific tests for CacheDiT enablement.
""" """
import ast import ast
import sys
from pathlib import Path from pathlib import Path
from unittest.mock import Mock, patch from unittest.mock import Mock, patch
@@ -25,18 +24,9 @@ from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import L
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
from vllm_omni.platforms import current_omni_platform from vllm_omni.platforms import current_omni_platform
# NOTE: We patch DreamID Omni's modules here with mocks so that we can import and inspect
# the class even though the dependency may not be set up correctly; this is ok for these
# tests because we just inspect it and never initialize the model.
for mod in ("dreamid_omni", "dreamid_omni.modules", "dreamid_omni.modules.model"):
sys.modules.setdefault(mod, Mock())
# isort: split
from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel as DreamIdOmniModel # noqa: E402
pytestmark = [pytest.mark.core_model, pytest.mark.cpu] pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
SEPARATE_CFG_TRANSFORMERS = [ SEPARATE_CFG_TRANSFORMERS = [
DreamIdOmniModel,
HeliosTransformer3DModel, HeliosTransformer3DModel,
LongCatImageTransformer2DModel, LongCatImageTransformer2DModel,
Cosmos3VFMTransformer, Cosmos3VFMTransformer,
@@ -9,6 +9,7 @@ Uses ``Tongyi-MAI/Z-Image-Turbo`` as the parity model.
from __future__ import annotations from __future__ import annotations
import gc
import os import os
import time import time
from pathlib import Path from pathlib import Path
@@ -75,8 +76,13 @@ def _run_generate(
if num_requests < 2: if num_requests < 2:
raise ValueError("num_requests must be >= 2 (1 warmup + >=1 timed)") raise ValueError("num_requests must be >= 2 (1 warmup + >=1 timed)")
gc.collect()
current_omni_platform.empty_cache() current_omni_platform.empty_cache()
device_index = current_omni_platform.current_device() device_index = current_omni_platform.current_device()
with current_omni_platform.device(device_index):
free_bytes, total_bytes = current_omni_platform.mem_get_info()
initial_used_mb = (total_bytes - free_bytes) / (1024**2)
monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02) monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
monitor.start() monitor.start()
try: try:
@@ -113,7 +119,10 @@ def _run_generate(
pass pass
median_time_s = float(np.median(per_request_times_s)) median_time_s = float(np.median(per_request_times_s))
peak_memory_mb = monitor.peak_used_mb # DeviceMemoryMonitor samples absolute device usage. Normalize by
# this run's starting usage so compiler and allocator caches retained
# by the preceding TP=1 run are not charged to the TP=2 run.
peak_memory_mb = max(0.0, monitor.peak_used_mb - initial_used_mb)
return _extract_single_image([last_output]), median_time_s, peak_memory_mb return _extract_single_image([last_output]), median_time_s, peak_memory_mb
finally: finally:
monitor.stop() monitor.stop()
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# ruff: noqa: N803, E741 # ruff: noqa: N803, E741
"""Layer-level correctness & performance test for MoT parallel linear layers. """Layer-level correctness & performance test for MoT parallel linear layers.
@@ -237,6 +240,53 @@ def _reference_forward(x, text_indices, vae_indices, text_linear, vae_linear):
return output return output
def _assert_gen_expert_parallel_metadata(layer):
"""The secondary expert must satisfy vLLM's parallel-linear contract."""
for attr in (
"tp_rank",
"tp_size",
"input_size",
"input_size_per_partition",
"output_size",
"output_size_per_partition",
):
assert getattr(layer.gen_exp, attr) == getattr(layer, attr)
def test_mot_gen_expert_parallel_metadata():
"""Keep secondary experts compatible with vLLM online quantizers."""
from vllm.model_executor.layers.quantization.online.fp8 import (
_is_tp_sharded,
)
with set_current_vllm_config(VllmConfig()):
layers = (
MoTQKVParallelLinear(
hidden_size=16,
head_size=4,
total_num_heads=4,
total_num_kv_heads=2,
bias=False,
vae_bias=False,
params_dtype=torch.bfloat16,
disable_tp=True,
),
MoTRowParallelLinear(
16,
8,
bias=False,
vae_bias=False,
input_is_parallel=True,
params_dtype=torch.bfloat16,
disable_tp=True,
),
)
for layer in layers:
_assert_gen_expert_parallel_metadata(layer)
assert _is_tp_sharded(layer.gen_exp) is False
def _check_and_report(ref: torch.Tensor, mot: torch.Tensor, tag: str): def _check_and_report(ref: torch.Tensor, mot: torch.Tensor, tag: str):
"""Compare outputs, print metrics, and assert correctness. """Compare outputs, print metrics, and assert correctness.
@@ -335,6 +385,8 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str, bias: bool
disable_tp=True, disable_tp=True,
).cuda() ).cuda()
_assert_gen_expert_parallel_metadata(mot_linear)
assert text_linear.output_size_per_partition == N, ( assert text_linear.output_size_per_partition == N, (
f"Expected output_size_per_partition={N}, " f"Expected output_size_per_partition={N}, "
f"got {text_linear.output_size_per_partition}. " f"got {text_linear.output_size_per_partition}. "
@@ -423,6 +475,8 @@ def test_mot_o_proj(
disable_tp=True, disable_tp=True,
).cuda() ).cuda()
_assert_gen_expert_parallel_metadata(mot_linear)
_sync_weights(text_linear, vae_linear, mot_linear) _sync_weights(text_linear, vae_linear, mot_linear)
text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE) text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE)
@@ -501,6 +555,13 @@ def test_mot_qkv_und_mode(K: int, N: int, dtype: str, bias: bool):
).cuda() ).cuda()
with torch.no_grad(): with torch.no_grad():
# vLLM linear layers leave weights uninitialized (torch.empty), so
# a freshly zeroed GPU page would make both outputs zero and the
# cosine-similarity gate degenerate. Seed deterministic values so
# und-mode is compared against a meaningful reference.
ref_linear.weight.normal_(mean=0.0, std=0.02)
if ref_linear.bias is not None:
ref_linear.bias.normal_(mean=0.0, std=0.02)
mot_linear.weight.copy_(ref_linear.weight) mot_linear.weight.copy_(ref_linear.weight)
if bias and ref_linear.bias is not None: if bias and ref_linear.bias is not None:
mot_linear.bias.copy_(ref_linear.bias) mot_linear.bias.copy_(ref_linear.bias)
@@ -549,6 +610,13 @@ def test_mot_row_und_mode(K: int, N: int, dtype: str, bias: bool):
).cuda() ).cuda()
with torch.no_grad(): with torch.no_grad():
# vLLM linear layers leave weights uninitialized (torch.empty), so
# a freshly zeroed GPU page would make both outputs zero and the
# cosine-similarity gate degenerate. Seed deterministic values so
# und-mode is compared against a meaningful reference.
ref_linear.weight.normal_(mean=0.0, std=0.02)
if ref_linear.bias is not None:
ref_linear.bias.normal_(mean=0.0, std=0.02)
mot_linear.weight.copy_(ref_linear.weight) mot_linear.weight.copy_(ref_linear.weight)
if bias and ref_linear.bias is not None: if bias and ref_linear.bias is not None:
mot_linear.bias.copy_(ref_linear.bias) mot_linear.bias.copy_(ref_linear.bias)
+91
View File
@@ -340,6 +340,97 @@ def test_rmsnorm_npu_receives_gamma_with_configured_dtype(monkeypatch: pytest.Mo
assert captured["epsilon"].item() == pytest.approx(1e-5) assert captured["epsilon"].item() == pytest.approx(1e-5)
# ── RMSNorm residual tests ──
def test_rmsnorm_matches_native_residual_contract():
"""The non-NPU path preserves separate residual-add and RMSNorm math."""
from vllm_omni.diffusion.layers.norm import RMSNorm
eps = 1e-6
norm = RMSNorm(4, eps=eps, dtype=torch.bfloat16)
norm.weight.data.copy_(torch.tensor([1.0, 0.5, 1.5, 2.0], dtype=torch.bfloat16))
x = torch.tensor([[1.0, -2.0, 3.0, -4.0]], dtype=torch.bfloat16)
residual = torch.tensor([[0.00390625, 0.0078125, -0.0078125, 0.015625]], dtype=torch.bfloat16)
expected_residual = residual + x
expected_fp32 = expected_residual.float()
expected_output = expected_fp32 * torch.rsqrt(expected_fp32.square().mean(-1, keepdim=True) + eps)
expected_output = (expected_output * norm.weight.float()).to(x.dtype)
output, updated_residual = norm.forward_native(x, residual)
torch.testing.assert_close(updated_residual, expected_residual, atol=0, rtol=0)
torch.testing.assert_close(output, expected_output, atol=0, rtol=0)
@pytest.mark.parametrize(
"forward_name",
["forward_cuda", "forward_hip", "forward_musa", "forward_xpu"],
)
def test_rmsnorm_non_npu_residual_delegates_to_native(
monkeypatch: pytest.MonkeyPatch,
forward_name: str,
):
"""Platform wrappers keep residual arithmetic in forward_native."""
from vllm_omni.diffusion.layers.norm import RMSNorm
norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16)
x = torch.randn(2, 4, dtype=torch.bfloat16)
residual = torch.randn_like(x)
expected_output = torch.randn_like(x)
expected_residual = torch.randn_like(x)
captured: dict[str, object] = {}
def forward_native(native_x: torch.Tensor, native_residual: torch.Tensor | None = None):
captured.update(x=native_x, residual=native_residual)
return expected_output, expected_residual
monkeypatch.setattr(norm, "forward_native", forward_native)
output, updated_residual = getattr(norm, forward_name)(x, residual)
assert output is expected_output
assert updated_residual is expected_residual
assert captured["x"] is x
assert captured["residual"] is residual
def test_rmsnorm_npu_residual_uses_torch_npu_fused_op(monkeypatch: pytest.MonkeyPatch):
"""NPU keeps the direct fused AddRMSNorm path instead of IR fallback."""
from vllm_omni.diffusion.layers.norm import RMSNorm
captured: dict[str, object] = {}
fused_output = torch.randn(2, 4, dtype=torch.bfloat16)
fused_residual = torch.randn(2, 4, dtype=torch.bfloat16)
def npu_add_rms_norm(
x: torch.Tensor,
residual: torch.Tensor,
gamma: torch.Tensor,
epsilon: float,
):
captured.update(x=x, residual=residual, gamma=gamma, epsilon=epsilon)
return fused_output, torch.empty(0), fused_residual
monkeypatch.setitem(
sys.modules,
"torch_npu",
types.SimpleNamespace(npu_add_rms_norm=npu_add_rms_norm),
)
norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16)
x = torch.randn(2, 4, dtype=torch.bfloat16)
residual = torch.randn(2, 4, dtype=torch.bfloat16)
output, updated_residual = norm.forward_npu(x, residual)
assert output is fused_output
assert updated_residual is fused_residual
assert captured["x"] is x
assert captured["residual"] is residual
assert captured["gamma"] is norm.weight
assert captured["epsilon"] == 1e-5
# ── RMSNorm compile-path regression tests ── # ── RMSNorm compile-path regression tests ──
+34
View File
@@ -0,0 +1,34 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from vllm_omni.diffusion.layers.norm import RMSNorm
from vllm_omni.platforms import current_omni_platform
pytestmark = [
pytest.mark.core_model,
pytest.mark.diffusion,
pytest.mark.cuda,
pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="CUDA platform required"),
]
def test_rmsnorm_cuda_residual_matches_native_contract() -> None:
eps = 1e-6
device = torch.device(current_omni_platform.device_type)
torch.manual_seed(0)
norm = RMSNorm(64, eps=eps, dtype=torch.bfloat16).to(device)
x = torch.randn(2, 4, 64, device=device, dtype=torch.bfloat16)
residual = torch.randn_like(x)
x_before = x.clone()
residual_before = residual.clone()
expected_output, expected_residual = norm.forward_native(x_before, residual_before)
output, updated_residual = norm.forward_cuda(x, residual)
torch.testing.assert_close(updated_residual, expected_residual, atol=0, rtol=0)
torch.testing.assert_close(output, expected_output, atol=0, rtol=0)
torch.testing.assert_close(x, x_before, atol=0, rtol=0)
torch.testing.assert_close(residual, residual_before, atol=0, rtol=0)
+24 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations from __future__ import annotations
@@ -89,6 +89,29 @@ def test_diffusion_base_linear_apply_multi_slice():
assert torch.allclose(out, expected) assert torch.allclose(out, expected)
def test_diffusion_lora_buffer_device_persists_across_reallocation(monkeypatch):
from vllm.lora.layers.base_linear import BaseLinearLayerWithLoRA
layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
torch.nn.Module.__init__(layer)
layer.base_layer = torch.nn.Module()
layer.n_slices = 1
layer.lora_a_stacked = (torch.zeros(2),)
layer.lora_b_stacked = (torch.zeros(3),)
layer._set_diffusion_lora_buffer_device(torch.device("meta"))
assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
def recreate_on_cpu(self, max_loras, lora_config, model_config):
self.lora_a_stacked = (torch.zeros(4),)
self.lora_b_stacked = (torch.zeros(5),)
monkeypatch.setattr(BaseLinearLayerWithLoRA, "create_lora_weights", recreate_on_cpu)
layer.create_lora_weights(max_loras=1, lora_config=object())
assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
def test_diffusion_base_linear_reset_lora_disables_fast_path(monkeypatch): def test_diffusion_base_linear_reset_lora_disables_fast_path(monkeypatch):
# Verify that after reset_lora(), apply() skips LoRA matmuls even if the # Verify that after reset_lora(), apply() skips LoRA matmuls even if the
# LoRA tensors are still allocated and non-empty. # LoRA tensors are still allocated and non-empty.
+33
View File
@@ -3,6 +3,8 @@
from __future__ import annotations from __future__ import annotations
from types import SimpleNamespace
import pytest import pytest
import torch import torch
from vllm.lora.lora_weights import LoRALayerWeights from vllm.lora.lora_weights import LoRALayerWeights
@@ -143,6 +145,37 @@ def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch):
assert replace_calls == ["foo"] assert replace_calls == ["foo"]
def test_lora_manager_keeps_dlo_lora_buffers_on_compute_device(monkeypatch):
import vllm_omni.diffusion.lora.manager as manager_mod
class _DLOLoRALayer(DummyBaseLayerWithLoRA):
lora_buffer_device = None
def _set_diffusion_lora_buffer_device(self, device):
self.lora_buffer_device = device
monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DLOLoRALayer)
monkeypatch.setattr(
manager_mod,
"from_layer_diffusion",
lambda *, layer, **_kwargs: _DLOLoRALayer(layer),
)
pipeline = torch.nn.Module()
pipeline.od_config = SimpleNamespace(enable_distributed_layerwise_offload=True)
pipeline.transformer = torch.nn.Module()
pipeline.transformer.foo = _FakeLinearBase()
manager = DiffusionLoRAManager(
pipeline=pipeline,
device=torch.device("cpu"),
dtype=torch.bfloat16,
)
manager._replace_layers_with_lora(type("_PH", (), {"r": 1})())
assert pipeline.transformer.foo.lora_buffer_device == torch.device("cpu")
def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch): def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch):
import vllm_omni.diffusion.lora.manager as manager_mod import vllm_omni.diffusion.lora.manager as manager_mod
@@ -249,6 +249,30 @@ def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_p
assert loader.take_host_weight_plan() is None assert loader.take_host_weight_plan() is None
def test_required_hwr_miss_fails_before_ordinary_loading_or_publication(
tmp_path: Path,
monkeypatch,
):
canonical_root = tmp_path / "canonical"
canonical_root.mkdir()
save_file(
{"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)},
str(canonical_root / "model.safetensors"),
)
loader = DiffusersPipelineLoader(
LoadConfig(),
_hwr_config(canonical_root, tmp_path / "empty-store", mode="required"),
)
pipeline = _HWRPipeline(canonical_root)
monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline)
with pytest.raises(RuntimeError, match="Host Weight Runtime resolution failed"):
loader.load_model(load_device="cpu", device=torch.device("cpu"))
assert pipeline.load_count == 0
assert loader.take_host_weight_plan() is None
def _make_dlo_online_quant_config() -> OmniDiffusionConfig: def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
return OmniDiffusionConfig( return OmniDiffusionConfig(
model="", model="",
@@ -261,10 +285,11 @@ def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
@pytest.mark.parametrize( @pytest.mark.parametrize(
("dist_offload", "use_allgather"), ("dist_offload", "use_allgather", "mode"),
[ [
(False, False), (False, False, "preferred"),
(True, True), (True, True, "preferred"),
(True, False, "disabled"),
], ],
) )
def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction( def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
@@ -272,12 +297,13 @@ def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
tmp_path, tmp_path,
dist_offload, dist_offload,
use_allgather, use_allgather,
mode,
): ):
"""Disabled and AllGather paths must never construct or probe HWR.""" """Disabled and AllGather paths must never construct or probe HWR."""
from vllm_omni.host_weight_runtime import HostWeightRuntime from vllm_omni.host_weight_runtime import HostWeightRuntime
root = tmp_path / "must-not-be-touched" root = tmp_path / "must-not-be-touched"
loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root)) loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root, mode=mode))
model = _DummyPipelineModel(source_prefix="transformer.") model = _DummyPipelineModel(source_prefix="transformer.")
modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,)) modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,))
@@ -303,6 +303,83 @@ class _FakeH3Attention(nn.Module):
super().__init__() super().__init__()
class _FakeFluxLinear(_FakeH3Linear):
def __init__(
self,
*args: object,
bias: bool = False,
params_dtype: torch.dtype | None = None,
total_num_heads: int | None = None,
total_num_kv_heads: int | None = None,
**kwargs: object,
) -> None:
super().__init__(
*args,
bias=bias,
params_dtype=params_dtype,
total_num_heads=total_num_heads,
total_num_kv_heads=total_num_kv_heads,
**kwargs,
)
self.weight.weight_loader = self._load_weight
@staticmethod
def _load_weight(
parameter: nn.Parameter,
loaded_weight: torch.Tensor,
shard_id: str | None = None,
) -> None:
del shard_id
parameter.data.copy_(loaded_weight.to(parameter).reshape_as(parameter))
class _FakeFluxAttention(nn.Module):
def __init__(self, **kwargs: object) -> None:
del kwargs
super().__init__()
self.register_buffer("beta", torch.tensor([0.5]))
self.register_buffer("eps", torch.tensor([-1e-6]))
class _Flux2KleinPipeline(nn.Module):
def __init__(self, transformer: nn.Module) -> None:
super().__init__()
self.transformer = transformer
self.text_encoder = nn.Linear(2, 2, dtype=torch.bfloat16)
self.vae = nn.Module()
self.vae.register_buffer("gain", torch.tensor([9.0], dtype=torch.float32))
def _reduced_flux2_klein(
monkeypatch: pytest.MonkeyPatch,
*,
device: torch.device | str = "cpu",
) -> nn.Module:
from vllm_omni.diffusion.models.flux2_klein import flux2_klein_transformer as flux2
monkeypatch.setattr(flux2, "ColumnParallelLinear", _FakeFluxLinear)
monkeypatch.setattr(flux2, "MergedColumnParallelLinear", _FakeFluxLinear)
monkeypatch.setattr(flux2, "QKVParallelLinear", _FakeFluxLinear)
monkeypatch.setattr(flux2, "RowParallelLinear", _FakeFluxLinear)
monkeypatch.setattr(flux2, "Attention", _FakeFluxAttention)
transformer = flux2.Flux2Transformer2DModel(
patch_size=1,
in_channels=8,
out_channels=8,
num_layers=1,
num_single_layers=1,
attention_head_dim=8,
num_attention_heads=1,
joint_attention_dim=8,
timestep_guidance_channels=8,
mlp_ratio=1.0,
axes_dims_rope=(2, 2, 2, 2),
guidance_embeds=False,
)
return transformer.to(device=device, dtype=torch.bfloat16)
class _H3Pipeline(nn.Module): class _H3Pipeline(nn.Module):
def __init__(self, transformer: nn.Module) -> None: def __init__(self, transformer: nn.Module) -> None:
super().__init__() super().__init__()
@@ -1200,3 +1277,152 @@ def test_minimax_h3_declares_the_final_layout_restore_contract() -> None:
assert contract.implementation_id == "minimax-h3-dit" assert contract.implementation_id == "minimax-h3-dit"
assert callable(MiniMaxH3DiTModel.validate_restored_host_weights) assert callable(MiniMaxH3DiTModel.validate_restored_host_weights)
assert "data_parallel" not in FinalLayoutRequest.__dataclass_fields__ assert "data_parallel" not in FinalLayoutRequest.__dataclass_fields__
def test_reduced_flux2_klein_declares_complete_final_layout_contract(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
Flux2Transformer2DModel,
)
transformer = _reduced_flux2_klein(monkeypatch)
pipeline = _Flux2KleinPipeline(transformer)
source = _prepared_source(tmp_path, directory="flux2-klein-source")
context = _identity(pipeline, source)
targets = collect_final_layout_targets(
pipeline,
(("transformer", transformer),),
policy=FINAL_LAYOUT_BF16_POLICY,
require_materialized=True,
)
target_names = {target.name for target in targets}
contract = Flux2Transformer2DModel.host_weight_restore_contract
assert isinstance(contract, FinalLayoutModelContract)
assert contract.schema == FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA
assert contract.implementation_id == "flux2-klein-dit"
assert contract.version == "1"
assert context.model_contracts == (contract,)
assert transformer.stacked_params_mapping == [
(".to_qkv.", ".to_q.", "q"),
(".to_qkv.", ".to_k.", "k"),
(".to_qkv.", ".to_v.", "v"),
(".add_kv_proj", ".add_q_proj", "q"),
(".add_kv_proj", ".add_k_proj", "k"),
(".add_kv_proj", ".add_v_proj", "v"),
]
loaded = transformer.load_weights(
[
("transformer_blocks.0.attn.to_q.weight", torch.tensor([1.0], dtype=torch.bfloat16)),
("transformer_blocks.0.attn.add_q_proj.weight", torch.tensor([2.0], dtype=torch.bfloat16)),
(
"single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
torch.tensor([3.0], dtype=torch.bfloat16),
),
]
)
assert loaded == {
"transformer_blocks.0.attn.to_qkv.weight",
"transformer_blocks.0.attn.add_kv_proj.weight",
"single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
}
assert any(name.startswith("transformer.transformer_blocks.0.") for name in target_names)
assert any(name.startswith("transformer.single_transformer_blocks.0.") for name in target_names)
assert "transformer.transformer_blocks.0.attn.to_qkv.weight" in target_names
assert "transformer.transformer_blocks.0.attn.add_kv_proj.weight" in target_names
assert "transformer.single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight" in target_names
assert "transformer.transformer_blocks.0.attn.attn.beta" in target_names
assert "transformer.transformer_blocks.0.attn.attn.eps" in target_names
assert all("text_encoder" not in name and "vae" not in name for name in target_names)
transformer.validate_restored_host_weights()
def test_flux2_klein_restore_validator_rejects_non_ready_state(
monkeypatch: pytest.MonkeyPatch,
) -> None:
transformer = _reduced_flux2_klein(monkeypatch)
mapping = transformer.stacked_params_mapping
transformer.stacked_params_mapping = None
with pytest.raises(ValueError, match="packed QKV mapping"):
transformer.validate_restored_host_weights()
transformer.stacked_params_mapping = mapping
weight = transformer.x_embedder.weight
weight.data = weight.data.float()
with pytest.raises(ValueError, match="must stay bf16"):
transformer.validate_restored_host_weights()
weight.data = weight.data.to(torch.bfloat16)
weight.data = weight.data.T
with pytest.raises(ValueError, match="contiguous strided layout"):
transformer.validate_restored_host_weights()
weight.data = weight.data.contiguous()
attention = transformer.transformer_blocks[0].attn.attn
attention._non_persistent_buffers_set.add("beta")
with pytest.raises(ValueError, match="must be persistent"):
transformer.validate_restored_host_weights()
attention._non_persistent_buffers_set.remove("beta")
attention.beta.data = attention.beta.data.float()
with pytest.raises(ValueError, match="scalar bf16"):
transformer.validate_restored_host_weights()
def test_flux2_klein_publication_restore_fallback_and_teardown(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
source = _prepared_source(tmp_path, directory="flux2-klein-source")
cold_transformer = _reduced_flux2_klein(monkeypatch)
cold_model = _Flux2KleinPipeline(cold_transformer)
with torch.no_grad():
for index, parameter in enumerate(cold_transformer.parameters(), start=1):
parameter.fill_(index)
for name, buffer in cold_transformer.named_buffers():
if name.endswith((".beta", ".eps")):
buffer.fill_(0.5 if name.endswith(".beta") else -1e-6)
context = _identity(cold_model, source)
runtime = _runtime(tmp_path / "flux2-klein-store")
miss = runtime.resolve(context.identity)
assert miss.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
producer = FinalLayoutBF16Producer(
context,
cold_model,
(("transformer", cold_transformer),),
max_shard_bytes=64,
)
publication = runtime.publish_after_load(context.identity, producer=producer)
assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED
warm_transformer = _reduced_flux2_klein(monkeypatch, device="meta")
warm_model = _Flux2KleinPipeline(warm_transformer)
warm_context = _identity(warm_model, source)
assert warm_context.identity == context.identity
hit = runtime.resolve(warm_context.identity)
assert hit.report.outcome is ResolutionOutcome.LOCAL_HIT
assert hit.lease is not None
plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease)
assert all(tensor.is_meta for name, tensor in warm_model.named_parameters() if name.startswith("transformer."))
plan.commit()
warm_transformer.validate_restored_host_weights()
cold_tensors = dict(cold_model.named_parameters()) | dict(cold_model.named_buffers())
warm_tensors = dict(warm_model.named_parameters()) | dict(warm_model.named_buffers())
for name in warm_context.tensor_names:
assert torch.equal(warm_tensors[name], cold_tensors[name])
assert warm_tensors[name].untyped_storage().data_ptr() == hit.lease.tensors[name].untyped_storage().data_ptr()
assert warm_transformer.stacked_params_mapping == cold_transformer.stacked_params_mapping
del warm_tensors, warm_model, warm_transformer, plan
gc.collect()
hit.lease.close()
assert hit.lease.closed
assert isinstance(runtime.store, FilesystemHostWeightStore)
assert runtime.store.cleanup(context.identity) is None
@@ -1,27 +0,0 @@
import pytest
import torch.nn as nn
pytest.importorskip("dreamid_omni.modules.model")
from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_dreamid_omni_exposes_hsdp_shard_conditions_for_fused_blocks():
model = object.__new__(FusionModel)
nn.Module.__init__(model)
model.fused_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
model.attn = nn.Identity()
conditions = getattr(model, "_hsdp_shard_conditions", None)
assert conditions is not None
assert len(conditions) == 1
matched = []
for name, module in model.named_modules():
if any(cond(name, module) for cond in conditions):
matched.append(name)
assert matched == ["fused_blocks.0", "fused_blocks.1"]
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations from __future__ import annotations
@@ -245,7 +245,7 @@ else:
[sys.executable, "-c", script, str(_MODULE_PATH), str(tmp_path)], [sys.executable, "-c", script, str(_MODULE_PATH), str(tmp_path)],
capture_output=True, capture_output=True,
text=True, text=True,
timeout=3, timeout=30,
check=False, check=False,
) )
@@ -1,11 +1,12 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Focused LTX-2.5 pipeline correctness tests.""" """Focused LTX-2.5 pipeline correctness tests."""
import json import json
from dataclasses import replace from dataclasses import replace
from types import SimpleNamespace from types import SimpleNamespace
from typing import Any
import pytest import pytest
import torch import torch
@@ -170,7 +171,10 @@ def _resolve_request_inputs_for_test(
def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatch): def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatch):
pipe = SimpleNamespace(component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None)) pipe = SimpleNamespace(
component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None),
use_diffusion_decoder=True,
)
od_config = SimpleNamespace(model="Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16) od_config = SimpleNamespace(model="Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16)
monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu")) monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
@@ -187,7 +191,19 @@ def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatc
def test_ltx_converted_component_loading_propagates_revision(monkeypatch): def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
revision = "pinned-revision" revision = "pinned-revision"
calls = {"components": []} calls: dict[str, Any] = {"components": []}
natten_processor = object()
class FakeDiffusionDecoder:
def set_attn_processor(self, processor):
calls["diffusion_decoder_processor"] = processor
def enable_tiling(self):
calls["diffusion_decoder_tiling"] = True
def set_parallel_size(self, parallel_size, mode):
calls["diffusion_decoder_parallel"] = (parallel_size, mode)
profile = replace( profile = replace(
LTX25_DISTILLED_COMPONENT_PROFILE, LTX25_DISTILLED_COMPONENT_PROFILE,
text_encoder_cls=object, text_encoder_cls=object,
@@ -196,12 +212,19 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
vocoder_fallback_cls=None, vocoder_fallback_cls=None,
scheduler_use_dynamic_shifting=False, scheduler_use_dynamic_shifting=False,
) )
pipeline = SimpleNamespace(component_profile=profile, pipeline_kind="distilled_two_stage") pipeline = SimpleNamespace(
component_profile=profile,
pipeline_kind="distilled_two_stage",
use_diffusion_decoder=True,
)
od_config = SimpleNamespace( od_config = SimpleNamespace(
model="org/converted-ltx25", model="org/converted-ltx25",
revision=revision, revision=revision,
dtype=torch.bfloat16, dtype=torch.bfloat16,
quantization_config=None, quantization_config=None,
extras={},
vae_use_tiling=False,
parallel_config=SimpleNamespace(vae_patch_parallel_size=2, vae_parallel_mode="tile"),
) )
monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu")) monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
@@ -227,6 +250,10 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
) )
return object() return object()
def fake_native_diffusion_decoder(model, **kwargs):
calls["native_diffusion_decoder"] = (model, kwargs)
return FakeDiffusionDecoder()
def fake_transformer_config(model, subfolder, local_files_only, *, revision): def fake_transformer_config(model, subfolder, local_files_only, *, revision):
calls["transformer_config"] = (model, subfolder, local_files_only, revision) calls["transformer_config"] = (model, subfolder, local_files_only, revision)
return {"component": "transformer"} return {"component": "transformer"}
@@ -234,7 +261,13 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
monkeypatch.setattr(ltx2_components, "prefetch_subfolders", fake_prefetch) monkeypatch.setattr(ltx2_components, "prefetch_subfolders", fake_prefetch)
monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", fake_tokenizer) monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", fake_tokenizer)
monkeypatch.setattr(ltx2_components, "_load_component", fake_component) monkeypatch.setattr(ltx2_components, "_load_component", fake_component)
monkeypatch.setattr(ltx2_components, "_load_ltx25_native_diffusion_decoder", fake_native_diffusion_decoder)
monkeypatch.setattr(ltx2_components, "load_transformer_config", fake_transformer_config) monkeypatch.setattr(ltx2_components, "load_transformer_config", fake_transformer_config)
monkeypatch.setattr(
ltx2_components,
"LTX2VideoVaeNeighborhoodNattenProcessor",
lambda: natten_processor,
)
monkeypatch.setattr( monkeypatch.setattr(
ltx2_components, ltx2_components,
"create_transformer_from_config", "create_transformer_from_config",
@@ -261,6 +294,14 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
"latent_upsampler", "latent_upsampler",
} }
assert all(call[3]["revision"] == revision for call in calls["components"]) assert all(call[3]["revision"] == revision for call in calls["components"])
assert "diffusion_decoder" not in calls["prefetch"][1]
assert calls["native_diffusion_decoder"] == (
od_config.model,
{"local_files_only": False, "dtype": torch.bfloat16, "revision": revision},
)
assert calls["diffusion_decoder_processor"] is natten_processor
assert calls["diffusion_decoder_tiling"] is True
assert calls["diffusion_decoder_parallel"] == (2, "tile")
assert calls["transformer_config"] == ( assert calls["transformer_config"] == (
od_config.model, od_config.model,
profile.transformer_subfolder, profile.transformer_subfolder,
@@ -270,6 +311,85 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
assert calls["scheduler"][1]["revision"] == revision assert calls["scheduler"][1]["revision"] == revision
def test_ltx25_native_diffusion_decoder_uses_diffusers_config_and_canonical_artifact(monkeypatch):
revision = "diffusers-revision"
config = {"decoder_stage_channels": [8, 8, 8, 8, 8]}
calls: dict[str, Any] = {}
class FakeDecoder:
def init_distributed(self):
calls["init_distributed"] = True
expected = FakeDecoder()
def fake_load_config(model, **kwargs):
calls["config"] = (model, kwargs)
return config
def fake_resolve(model, repo_id, filename, **kwargs):
calls["artifact"] = (model, repo_id, filename, kwargs)
return "/models/native-diffvae.safetensors"
def fake_load(cls, path, passed_config, dtype):
calls["load"] = (path, passed_config, dtype)
return expected
monkeypatch.setattr(
ltx2_components.DistributedLTX2VideoDiffusionDecoderModel,
"load_config",
staticmethod(fake_load_config),
)
monkeypatch.setattr(ltx2_components, "resolve_ltx_artifact", fake_resolve)
monkeypatch.setattr(
ltx2_components.DistributedLTX2VideoDiffusionDecoderModel,
"from_ltx25_native_checkpoint",
classmethod(fake_load),
)
actual = ltx2_components._load_ltx25_native_diffusion_decoder(
"Lightricks/LTX-2.5-Diffusers",
local_files_only=False,
dtype=torch.bfloat16,
revision=revision,
)
assert actual is expected
assert calls["config"] == (
"Lightricks/LTX-2.5-Diffusers",
{
"subfolder": "diffusion_decoder",
"local_files_only": False,
"revision": revision,
},
)
assert calls["artifact"] == (
"Lightricks/LTX-2.5-Diffusers",
"Lightricks/LTX-2.5",
"vae/ltx-2.5-video-vae-bf16.safetensors",
{
"model_revision": revision,
"artifact_revision": ltx2_components.LTX25_NATIVE_ARTIFACT_REVISION,
},
)
assert calls["load"] == ("/models/native-diffvae.safetensors", config, torch.bfloat16)
assert calls["init_distributed"] is True
def test_ltx25_natten_failure_has_actionable_omni_remedy(monkeypatch):
class MissingNattenProcessor:
def __init__(self):
raise ImportError("kernels is unavailable")
monkeypatch.setattr(
ltx2_components,
"LTX2VideoVaeNeighborhoodNattenProcessor",
MissingNattenProcessor,
)
with pytest.raises(RuntimeError, match=r"kernels==0\.14\.1.*supported GPU.*allow Hub access"):
ltx2_components._create_ltx25_natten_processor()
def test_ltx_checkpoint_explicit_version_precedes_structural_heuristics(tmp_path): def test_ltx_checkpoint_explicit_version_precedes_structural_heuristics(tmp_path):
(tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"})) (tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"}))
(tmp_path / "transformer").mkdir() (tmp_path / "transformer").mkdir()
@@ -411,6 +531,46 @@ def test_ltx25_four_public_pipeline_semantics_are_disjoint():
assert LTX2DistilledTwoStagePipeline.pipeline_kind == "distilled_two_stage" assert LTX2DistilledTwoStagePipeline.pipeline_kind == "distilled_two_stage"
@pytest.mark.parametrize(
"pipeline_cls",
[
LTX2Pipeline,
LTX2TwoStagePipeline,
LTX2DistilledOneStagePipeline,
LTX2DistilledTwoStagePipeline,
],
)
@pytest.mark.parametrize(
("extras", "use_diffusion_decoder"),
[({}, True), ({"ltx2_use_conv_vae": True}, False)],
)
def test_ltx25_all_public_pipelines_default_to_diffvae_with_conv_opt_in(
tmp_path, monkeypatch, pipeline_cls, extras, use_diffusion_decoder
):
(tmp_path / "model_index.json").write_text(
json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
)
def stub_components(pipe, od_config):
pipe.od_config = od_config
pipe.vae_spatial_compression_ratio = 32
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda _pipe: None)
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
pipe = pipeline_cls(
od_config=SimpleNamespace(
model=str(tmp_path),
extras=extras,
enable_diffusion_pipeline_profiler=False,
)
)
assert pipe.use_diffusion_decoder is use_diffusion_decoder
assert pipe._vae_modules.count("diffusion_decoder") == int(use_diffusion_decoder)
@pytest.mark.parametrize( @pytest.mark.parametrize(
("pipeline_kind", "expected"), ("pipeline_kind", "expected"),
[ [
@@ -673,7 +833,7 @@ def test_ltx25_distilled_two_stage_executes_custom_phase_schedules():
max_sequence_length=16, max_sequence_length=16,
) )
prompt_context = object() prompt_context = object()
phase_calls = [] phase_calls: list[Any] = []
def resolve_request_inputs(req, **kwargs): def resolve_request_inputs(req, **kwargs):
return request_inputs return request_inputs
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for LTX image-to-video input and conditioning behavior.""" """Unit tests for LTX image-to-video input and conditioning behavior."""
@@ -258,6 +258,39 @@ class TestLTXImageToVideoConditioning:
assert LTX2Pipeline.support_image_input is True assert LTX2Pipeline.support_image_input is True
def test_ltx25_i2v_encode_uses_diffusion_decoder_statistics(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.use_diffusion_decoder = True
pipe.vae = SimpleNamespace(
encode=lambda image: image,
latents_mean=torch.full((2,), 100.0),
latents_std=torch.full((2,), 10.0),
config=SimpleNamespace(scaling_factor=1.0),
)
pipe.diffusion_decoder = SimpleNamespace(
latents_mean=torch.tensor([1.0, 3.0]),
latents_std=torch.tensor([2.0, 4.0]),
config=SimpleNamespace(scaling_factor=2.0),
)
monkeypatch.setattr(
ltx2_conditioning,
"retrieve_latents",
lambda *_args, **_kwargs: torch.tensor([[[[[3.0]]], [[[7.0]]]]]),
)
actual = pipe._encode_i2v_image_latents(
torch.zeros(1, 3, 1, 1),
batch_size=1,
generator=None,
dtype=torch.float32,
)
torch.testing.assert_close(actual, torch.full((1, 2, 1, 1, 1), 2.0))
def test_ltx23_i2v_rejects_multi_image_prompt_list(self): def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace from types import SimpleNamespace
import pytest import pytest
@@ -47,6 +50,39 @@ def test_prepare_video_latents_matches_official_values_and_token_major_layout():
assert actual.stride()[1:] == (1, actual.shape[1]) assert actual.stride()[1:] == (1, actual.shape[1])
def test_prepare_video_latents_uses_diffusion_decoder_statistics():
pipeline = _make_pipeline(LTX2Pipeline)
pipeline.use_diffusion_decoder = True
pipeline.vae_spatial_compression_ratio = 1
pipeline.vae_temporal_compression_ratio = 1
pipeline.transformer_spatial_patch_size = 1
pipeline.transformer_temporal_patch_size = 1
pipeline.vae = SimpleNamespace(
latents_mean=torch.full((2,), 100.0),
latents_std=torch.full((2,), 10.0),
config=SimpleNamespace(scaling_factor=1.0),
)
pipeline.diffusion_decoder = SimpleNamespace(
latents_mean=torch.tensor([1.0, 3.0]),
latents_std=torch.tensor([2.0, 4.0]),
config=SimpleNamespace(scaling_factor=2.0),
)
actual = pipeline.prepare_latents(
batch_size=1,
num_channels_latents=2,
height=1,
width=1,
num_frames=1,
noise_scale=0.0,
dtype=torch.float32,
device=torch.device("cpu"),
latents=torch.tensor([[[[[3.0]]], [[[7.0]]]]]),
)
torch.testing.assert_close(actual, torch.full((1, 1, 2), 2.0))
def test_prepare_audio_latents_samples_directly_in_packed_token_space(): def test_prepare_audio_latents_samples_directly_in_packed_token_space():
pipeline = _make_pipeline(LTX2Pipeline) pipeline = _make_pipeline(LTX2Pipeline)
@@ -1,9 +1,10 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for the fixed LTX refinement-phase adapter.""" """Unit tests for the fixed LTX refinement-phase adapter."""
from types import SimpleNamespace from types import SimpleNamespace
from typing import Any
import pytest import pytest
import torch import torch
@@ -148,10 +149,13 @@ def test_ltx_phase_adapter_selects_execution_from_quantization(
): ):
pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config) pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config)
manifest = object() manifest = object()
captured = {} captured: dict[str, Any] = {}
def resolve_artifact(*_args, **kwargs): def resolve_artifact(*_args, **kwargs):
assert not kwargs assert kwargs == {
"model_revision": None,
"artifact_revision": None,
}
return "adapter" return "adapter"
def parse(transformer, path): def parse(transformer, path):
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for the shared LTX pipeline runtime and public contracts.""" """Unit tests for the shared LTX pipeline runtime and public contracts."""
@@ -8,6 +8,7 @@ import os
import tempfile import tempfile
from dataclasses import replace from dataclasses import replace
from types import SimpleNamespace from types import SimpleNamespace
from typing import Any
import pytest import pytest
import torch import torch
@@ -165,7 +166,19 @@ def test_ltx_checkpoint_version_detection_uses_metadata(tmp_path):
assert detect_ltx_model_version(str(tmp_path)) == "2" assert detect_ltx_model_version(str(tmp_path)) == "2"
def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch): @pytest.mark.parametrize(
("model", "repo_id", "expected_revision"),
[
("Lightricks/test", "Lightricks/test", "model-revision"),
("Lightricks/test-Diffusers", "Lightricks/test", "artifact-revision"),
],
)
def test_ltx_artifact_uses_source_revision_and_hub_fallback(
monkeypatch,
model,
repo_id,
expected_revision,
):
filename = "ltx-sidecar.safetensors" filename = "ltx-sidecar.safetensors"
calls = [] calls = []
@@ -176,9 +189,22 @@ def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch)
monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download) monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
assert ( assert (
resolve_ltx_artifact(str(tmp_path / "model"), "Lightricks/test", filename) == "/cache/ltx-sidecar.safetensors" resolve_ltx_artifact(
model,
repo_id,
filename,
model_revision="model-revision",
artifact_revision="artifact-revision",
)
== "/cache/ltx-sidecar.safetensors"
) )
assert calls == [{"repo_id": "Lightricks/test", "filename": filename}] assert calls == [
{
"repo_id": repo_id,
"filename": filename,
"revision": expected_revision,
}
]
def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch): def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
@@ -187,7 +213,41 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
expected.write_bytes(b"sidecar") expected.write_bytes(b"sidecar")
monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup")) monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup"))
assert resolve_ltx_artifact(str(tmp_path), "Lightricks/test", filename) == str(expected) assert resolve_ltx_artifact(
str(tmp_path),
"Lightricks/test",
filename,
model_revision="unused-model-revision",
artifact_revision="unused-artifact-revision",
) == str(expected)
def test_ltx_artifact_local_model_missing_sidecar_falls_back_to_hub(tmp_path, monkeypatch):
calls = []
def fake_download(**kwargs):
calls.append(kwargs)
return "/cache/ltx-sidecar.safetensors"
monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
assert (
resolve_ltx_artifact(
str(tmp_path),
"Lightricks/test",
"ltx-sidecar.safetensors",
model_revision="local-model-revision",
artifact_revision="artifact-revision",
)
== "/cache/ltx-sidecar.safetensors"
)
assert calls == [
{
"repo_id": "Lightricks/test",
"filename": "ltx-sidecar.safetensors",
"revision": "artifact-revision",
}
]
@pytest.mark.parametrize( @pytest.mark.parametrize(
@@ -900,7 +960,7 @@ def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
max_sequence_length=16, max_sequence_length=16,
) )
prompt_context = object() prompt_context = object()
phase_calls = [] phase_calls: list[Any] = []
def resolve_request_inputs(req, **kwargs): def resolve_request_inputs(req, **kwargs):
return request_inputs return request_inputs
@@ -959,7 +1019,7 @@ def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
pipeline.vae_spatial_compression_ratio = 32 pipeline.vae_spatial_compression_ratio = 32
pipeline.vae_temporal_compression_ratio = 8 pipeline.vae_temporal_compression_ratio = 8
pipeline.latent_upsampler = FakeUpsampler() pipeline.latent_upsampler = FakeUpsampler()
activated_slots = [] activated_slots: list[Any] = []
if pipeline_cls is LTX2TwoStagePipeline: if pipeline_cls is LTX2TwoStagePipeline:
pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append) pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append)
object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs) object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
@@ -1133,6 +1193,59 @@ def test_ltx_two_stage_recipe_rejects_sigmas():
) )
def test_ltx_request_applies_diffvae_minimum_spatial_size_only_when_requested():
request = LTXRequestInputs(
prompt="prompt",
negative_prompt="",
height=128,
width=128,
num_frames=9,
frame_rate=24.0,
num_inference_steps=30,
guidance=LTX2_ONE_STAGE_RECIPE.request_guidance,
num_videos_per_prompt=1,
generator=None,
latents=None,
audio_latents=None,
prompt_embeds=None,
negative_prompt_embeds=None,
prompt_attention_mask=None,
negative_prompt_attention_mask=None,
decode_timestep=0.0,
decode_noise_scale=None,
output_type="np",
max_sequence_length=16,
)
with pytest.raises(ValueError, match="DiffVAE requires video latent spatial dimensions of at least 7x7"):
validate_pipeline_request(
request,
pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
vae_spatial_compression_ratio=32,
vae_temporal_compression_ratio=8,
pipeline_name="LTX2Pipeline",
min_video_latent_spatial_size=(7, 7),
)
# ConvVAE does not use NATTEN, so its existing 128x128 boundary remains valid.
validate_pipeline_request(
request,
pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
vae_spatial_compression_ratio=32,
vae_temporal_compression_ratio=8,
pipeline_name="LTX2Pipeline",
)
validate_pipeline_request(
replace(request, height=256, width=256),
pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
vae_spatial_compression_ratio=32,
vae_temporal_compression_ratio=8,
pipeline_name="LTX2Pipeline",
min_video_latent_spatial_size=(7, 7),
)
@pytest.mark.parametrize( @pytest.mark.parametrize(
("direct_kwargs", "sampling_kwargs", "error"), ("direct_kwargs", "sampling_kwargs", "error"),
[ [
@@ -1284,7 +1397,7 @@ def test_ltx_official_two_stage_recipes_only_vary_by_model_defaults(recipe, step
def test_ltx_two_stage_entries_select_the_official_adapter_slots(): def test_ltx_two_stage_entries_select_the_official_adapter_slots():
phase_switches = [] phase_switches: list[Any] = []
ordinary_pipeline = object.__new__(LTX2TwoStagePipeline) ordinary_pipeline = object.__new__(LTX2TwoStagePipeline)
ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append) ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append)
for phase in LTX2_TWO_STAGE_RECIPE.phases: for phase in LTX2_TWO_STAGE_RECIPE.phases:
+464 -3
View File
@@ -1,9 +1,10 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for LTX-2.3 VAE tiling and distributed decode behavior.""" """Unit tests for LTX video VAE tiling and distributed decode behavior."""
from types import SimpleNamespace from types import SimpleNamespace
from typing import Any
import pytest import pytest
import torch import torch
@@ -11,6 +12,466 @@ import torch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu] pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
class FakeBaseVocoder(torch.nn.Module):
def __init__(self):
super().__init__()
self.input_dtype = None
def forward(self, value):
self.input_dtype = value.dtype
return value
monkeypatch.setattr(
ltx_runtime.torch,
"autocast",
lambda **_kwargs: (_ for _ in ()).throw(AssertionError("base vocoder must not use autocast")),
)
vocoder = FakeBaseVocoder()
output = ltx_runtime._run_ltx_vocoder(vocoder, torch.ones(1, dtype=torch.bfloat16))
assert vocoder.input_dtype == torch.bfloat16
assert output.dtype == torch.bfloat16
class TestLTXDiffusionDecoder:
@pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width"])
def test_distributed_diffusion_decoder_rejects_non_tile_parallel_modes(self, mode):
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
DistributedLTX2VideoDiffusionDecoderModel,
)
model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel)
torch.nn.Module.__init__(model)
model.distributed_executor = SimpleNamespace(set_parallel_size=lambda *_args, **_kwargs: None)
with pytest.raises(ValueError, match="only supports vae_parallel_mode='tile'"):
model.set_parallel_size(2, mode=mode)
def test_distributed_diffusion_decoder_accepts_tile_parallel_mode(self):
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
DistributedLTX2VideoDiffusionDecoderModel,
)
calls = []
model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel)
torch.nn.Module.__init__(model)
model.distributed_executor = SimpleNamespace(
set_parallel_size=lambda parallel_size, mode: calls.append((parallel_size, mode))
)
model.set_parallel_size(4, mode="tile")
assert calls == [(4, "tile")]
def test_short_clip_keeps_stage5_temporal_context_then_crops_output(self):
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
LTX2VideoDiffusionDecoder3d,
LTX2VideoDiffusionDecoderModel,
)
class IdentityAttention:
def build_block_mask(self, _hidden_states):
return None
class IdentityBlock(torch.nn.Module):
def __init__(self):
super().__init__()
self.attn = IdentityAttention()
def forward(self, hidden_states, _block_mask):
return hidden_states
class IdentityUpsample(torch.nn.Module):
def forward(self, hidden_states, *, drop_leading_frame):
assert drop_leading_frame
return hidden_states
decoder = object.__new__(LTX2VideoDiffusionDecoder3d)
torch.nn.Module.__init__(decoder)
decoder.det_stages = torch.nn.ModuleList([torch.nn.ModuleList([IdentityBlock()])])
decoder.upsamples = torch.nn.ModuleList([IdentityUpsample()])
decoder.trailing_pad_latent_frames = 2
decoder.temporal_compression_ratio = 8
decoder.stage5_kernel = (11, 3, 3)
# 17 frames contain one real frame followed by 16 propagated ghost
# frames. The reference keeps 11 for stage-5 NATTEN instead of
# cropping the tensor immediately back to one.
context = decoder.forward_stage_4(torch.zeros(1, 17, 3, 3, 4))
assert context.shape == (1, 11, 3, 3, 4)
class FakeDecoder(torch.nn.Module):
def forward(self, z, *, generator, num_inference_steps):
assert generator is None
assert num_inference_steps is None
return torch.zeros(z.shape[0], 3, 11, z.shape[3] * 2, z.shape[4] * 2)
model = object.__new__(LTX2VideoDiffusionDecoderModel)
torch.nn.Module.__init__(model)
model.decoder = FakeDecoder()
model.use_tiling = False
model.tile_sample_min_num_frames = 80
model.tile_sample_min_height = 768
model.tile_sample_min_width = 768
model.temporal_compression_ratio = 8
model.spatial_compression_ratio = 2
output = model.decode(torch.zeros(1, 4, 1, 2, 3), generator=None, return_dict=False)[0]
assert output.shape == (1, 3, 1, 4, 6)
@pytest.mark.parametrize(
("model_version", "extras", "expected"),
[
("2.5", {}, True),
("2.5", {"ltx2_use_conv_vae": True}, False),
("2.3", {}, False),
],
)
def test_ltx25_diffusion_decoder_is_default_with_conv_vae_opt_in(self, model_version, extras, expected):
from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder
assert _ltx2_use_diffusion_decoder(SimpleNamespace(extras=extras), model_version) is expected
def test_conv_vae_opt_in_rejects_non_boolean_model_extra(self):
from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder
with pytest.raises(TypeError, match="ltx2_use_conv_vae"):
_ltx2_use_diffusion_decoder(SimpleNamespace(extras={"ltx2_use_conv_vae": "true"}), "2.5")
def test_native_diffusion_decoder_conversion_splits_qkv_and_folds_gates(self):
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
convert_ltx25_native_diffusion_decoder_state_dict,
)
qkv_weight = torch.arange(24, dtype=torch.float32).reshape(6, 4)
qkv_bias = torch.arange(6, dtype=torch.float32)
projection_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2)
projection_bias = torch.ones(2)
gate = torch.tensor([2.0, 3.0])
mean = torch.tensor([0.1, 0.2])
std = torch.tensor([0.3, 0.4])
native = {
"encoder.ignored": torch.ones(1),
"per_channel_statistics.mean-of-means": mean,
"vae.per_channel_statistics.std-of-means": std,
"decoder.type_emb": torch.ones(2),
"decoder.coarse_head.weight": torch.ones(1),
"decoder.det_stages.0.0.attn.qkv.weight": qkv_weight,
"decoder.det_stages.0.0.attn.qkv.bias": qkv_bias,
"decoder.det_stages.0.0.attn.q_norm.weight": torch.tensor([5.0, 6.0]),
"decoder.det_stages.0.0.attn.k_norm.weight": torch.tensor([7.0, 8.0]),
"decoder.diff_blocks.0.attn.proj.weight": projection_weight,
"decoder.diff_blocks.0.attn.proj.bias": projection_bias,
"decoder.diff_blocks.0.gate_msa": gate,
"vae.decoder.t_embedder.mlp.0.weight": torch.ones(2, 2),
}
converted = convert_ltx25_native_diffusion_decoder_state_dict(native)
torch.testing.assert_close(converted["latents_mean"], mean)
torch.testing.assert_close(converted["latents_std"], std)
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.weight"], qkv_weight[:2])
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.weight"], qkv_weight[2:4])
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.weight"], qkv_weight[4:])
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.bias"], qkv_bias[:2])
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.bias"], qkv_bias[2:4])
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.bias"], qkv_bias[4:])
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_q.weight"], torch.tensor([5.0, 6.0]))
torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_k.weight"], torch.tensor([7.0, 8.0]))
torch.testing.assert_close(
converted["decoder.diff_blocks.0.attn.to_out.0.weight"], gate.unsqueeze(1) * projection_weight
)
torch.testing.assert_close(converted["decoder.diff_blocks.0.attn.to_out.0.bias"], gate * projection_bias)
assert "decoder.t_embedder.timestep_embedder.linear_1.weight" in converted
assert not any("type_emb" in key or "coarse" in key or "gate_msa" in key for key in converted)
def test_native_diffusion_decoder_conversion_rejects_invalid_fused_qkv(self):
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
convert_ltx25_native_diffusion_decoder_state_dict,
)
with pytest.raises(ValueError, match="not divisible by 3"):
convert_ltx25_native_diffusion_decoder_state_dict(
{"decoder.det_stages.0.0.attn.qkv.weight": torch.ones(5, 2)}
)
def test_decode_uses_diffusion_decoder_without_conv_vae_conditioning(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
seen = {}
class ConvVae:
dtype = torch.float32
config = SimpleNamespace(timestep_conditioning=True)
def decode(self, *_args, **_kwargs):
raise AssertionError("the convolutional VAE decoder must not run")
class DiffusionDecoder:
dtype = torch.float32
def decode(self, latents, *, generator, return_dict):
seen["latents"] = latents
seen["generator"] = generator
seen["return_dict"] = return_dict
return (latents + 1,)
class AudioVae:
dtype = torch.float32
def decode(self, audio_latents, *, return_dict):
assert return_dict is False
return (audio_latents + 2,)
class VideoProcessor:
def postprocess_video(self, video, *, output_type):
seen["output_type"] = output_type
return video
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.use_diffusion_decoder = True
pipe.vae = ConvVae()
pipe.diffusion_decoder = DiffusionDecoder()
pipe.audio_vae = AudioVae()
pipe.vocoder = torch.nn.Identity()
pipe.video_processor = VideoProcessor()
generator = torch.Generator().manual_seed(123)
latents = torch.ones(1, 1)
output = pipe._decode_output(
latents=latents,
audio_latents=torch.ones(1, 1),
output_type="pt",
connector_prompt_embeds=torch.ones(1, 1),
generator=generator,
device=torch.device("cpu"),
decode_timestep=0.7,
decode_noise_scale=0.5,
prompt_batch_size=1,
)
torch.testing.assert_close(seen["latents"], latents)
assert seen["generator"] is generator
assert seen["return_dict"] is False
assert seen["output_type"] == "pt"
torch.testing.assert_close(output.output[0], latents + 1)
torch.testing.assert_close(output.output[1], torch.full((1, 1), 3.0))
def test_diffusion_decode_denormalizes_with_decoder_statistics(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.use_diffusion_decoder = True
pipe.transformer_spatial_patch_size = 1
pipe.transformer_temporal_patch_size = 1
pipe.vae = SimpleNamespace(
latents_mean=torch.full((2,), 100.0),
latents_std=torch.full((2,), 10.0),
config=SimpleNamespace(scaling_factor=1.0),
)
pipe.diffusion_decoder = SimpleNamespace(
latents_mean=torch.tensor([10.0, 20.0]),
latents_std=torch.tensor([2.0, 4.0]),
config=SimpleNamespace(scaling_factor=2.0),
)
pipe.audio_vae = SimpleNamespace(
latents_mean=torch.tensor(0.0),
latents_std=torch.tensor(1.0),
)
forward_ctx = SimpleNamespace(
latent_num_frames=1,
latent_height=1,
latent_width=1,
original_audio_num_frames=1,
latent_mel_bins=2,
)
video, audio = pipe._unpack_and_denormalize_stage(
forward_ctx,
torch.tensor([[[1.0, 2.0]]]),
torch.zeros(1, 1, 2),
)
torch.testing.assert_close(video.flatten(), torch.tensor([11.0, 24.0]))
torch.testing.assert_close(audio, torch.zeros(1, 1, 1, 2))
def test_non_output_rank_does_not_enter_diffusion_decoder(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
class ConvVae:
config = SimpleNamespace(timestep_conditioning=False)
def is_distributed_enabled(self):
raise AssertionError("DiffVAE decode must not inspect ConvVAE collectives")
class DiffusionDecoder:
dtype = torch.float32
def decode(self, *_args, **_kwargs):
raise AssertionError("only the output rank may run DiffVAE decode")
monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True)
monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1)
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.use_diffusion_decoder = True
pipe.vae = ConvVae()
pipe.diffusion_decoder = DiffusionDecoder()
output = pipe._decode_output(
latents=torch.ones(1, 1),
audio_latents=torch.ones(1, 1),
output_type="pt",
connector_prompt_embeds=torch.ones(1, 1),
generator=None,
device=torch.device("cpu"),
decode_timestep=0.0,
decode_noise_scale=None,
prompt_batch_size=1,
)
assert output.output[0].numel() == 0
assert output.output[1].numel() == 0
def test_non_output_rank_enters_distributed_diffusion_decoder(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
class ConvVae:
config = SimpleNamespace(timestep_conditioning=False)
def is_distributed_enabled(self):
raise AssertionError("DiffVAE decode must query the active decoder")
class DiffusionDecoder:
dtype = torch.float32
def __init__(self):
self.decode_calls = 0
def is_distributed_enabled(self):
return True
def decode(self, *_args, **_kwargs):
self.decode_calls += 1
return (torch.empty(0, 3, 0, 0, 0),)
monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True)
monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1)
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.use_diffusion_decoder = True
pipe.vae = ConvVae()
pipe.diffusion_decoder = DiffusionDecoder()
output = pipe._decode_output(
latents=torch.ones(1, 1),
audio_latents=torch.ones(1, 1),
output_type="pt",
connector_prompt_embeds=torch.ones(1, 1),
generator=None,
device=torch.device("cpu"),
decode_timestep=0.0,
decode_noise_scale=None,
prompt_batch_size=1,
)
assert pipe.diffusion_decoder.decode_calls == 1
assert output.output[0].numel() == 0
assert output.output[1].numel() == 0
def test_diffusion_decoder_patch_parallel_size_one_uses_native_tiling(self, monkeypatch):
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
DistributedLTX2VideoDiffusionDecoderModel,
)
expected = torch.ones(1, 3, 1, 2, 2)
seen = {}
def native_tiled_decode(self, z, generator=None, num_inference_steps=None):
seen["z"] = z
seen["generator"] = generator
seen["num_inference_steps"] = num_inference_steps
return expected
monkeypatch.setattr(LTX2VideoDiffusionDecoderModel, "tiled_decode", native_tiled_decode)
model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel)
torch.nn.Module.__init__(model)
model.is_distributed_enabled = lambda: False
generator = torch.Generator().manual_seed(17)
z = torch.zeros(1, 1, 1, 1, 1)
result = DistributedLTX2VideoDiffusionDecoderModel.tiled_decode(
model,
z,
generator=generator,
num_inference_steps=1,
)
assert result is expected
assert seen["z"] is z
assert seen["generator"] is generator
assert seen["num_inference_steps"] == 1
def test_distributed_diffusion_tiles_preserve_serial_noise_order(self, monkeypatch):
from diffusers.utils.torch_utils import randn_tensor
from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
DistributedLTX2VideoDiffusionDecoderModel,
)
model = DistributedLTX2VideoDiffusionDecoderModel(
out_channels=1,
latent_channels=1,
patch_size=1,
decoder_head_dim=8,
decoder_stage_channels=(8, 8, 8, 8, 8),
decoder_stage_depths=(1, 1, 1, 1, 1),
decoder_stage_kernels=((1, 1, 1),) * 4,
decoder_upsample_strides=((1, 1, 1),) * 4,
decoder_upsample_channel_reductions=(1, 1, 1, 1),
decoder_stage5_kernel=(1, 1, 1),
spatial_compression_ratio=1,
temporal_compression_ratio=1,
)
model.tile_sample_min_num_frames = 2
model.tile_sample_stride_num_frames = 1
model.tile_sample_min_height = 2
model.tile_sample_stride_height = 1
model.tile_sample_min_width = 2
model.tile_sample_stride_width = 1
model.distributed_executor = SimpleNamespace(group=None)
model.decoder.forward_stages_1_to_3 = lambda z: z.permute(0, 2, 3, 4, 1)
monkeypatch.setattr(torch.distributed, "broadcast", lambda *_args, **_kwargs: None)
z = torch.zeros(1, 1, 2, 3, 3)
distributed_generator = torch.Generator().manual_seed(123)
reference_generator = torch.Generator().manual_seed(123)
tasks, grid_spec = model._distributed_tile_split(z, distributed_generator, num_inference_steps=1)
assert grid_spec.grid_shape == (1, 3, 3)
assert grid_spec.split_dims == (2, 3, 4)
for task in tasks:
tile_shape = model._tiled_pixel_shape_from_features(
task.tensor,
drop_leading_frame=task.drop_leading_frame,
crop_trailing_ghost=task.crop_trailing_ghost,
)
expected = randn_tensor(tile_shape, generator=reference_generator, device=z.device, dtype=z.dtype)
actual = randn_tensor(tile_shape, generator=task.noise_generator, device=z.device, dtype=z.dtype)
torch.testing.assert_close(actual, expected)
assert torch.equal(distributed_generator.get_state(), reference_generator.get_state())
class TestLTXOutputRank: class TestLTXOutputRank:
@pytest.mark.parametrize( @pytest.mark.parametrize(
("distributed_vae_state", "expected_decode_calls"), ("distributed_vae_state", "expected_decode_calls"),
@@ -113,7 +574,7 @@ class TestLTX23VaeDistributedDecode:
}, },
output_dtype=torch.float32, output_dtype=torch.float32,
) )
seen = {} seen: dict[str, Any] = {}
def exec_tile(task): def exec_tile(task):
return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1)) return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1))
@@ -0,0 +1,48 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""CUDA regression tests for LTX-2 BWE vocoder precision."""
import pytest
import torch
from tests.helpers.mark import hardware_test
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
class _BWEConvVocoder(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv1d(
1,
1,
kernel_size=1,
bias=False,
device="cuda",
dtype=torch.bfloat16,
)
self.bwe_generator = torch.nn.Identity()
self.input_dtype = None
self.conv_output_dtype = None
def forward(self, value):
self.input_dtype = value.dtype
output = self.conv(value)
self.conv_output_dtype = output.dtype
return output
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_ltx_bwe_vocoder_runs_real_cuda_autocast_in_fp32():
from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _run_ltx_vocoder
vocoder = _BWEConvVocoder()
generated_mel = torch.ones((1, 1, 4), device="cuda", dtype=torch.bfloat16)
output = _run_ltx_vocoder(vocoder, generated_mel)
assert next(vocoder.parameters()).dtype == torch.bfloat16
assert vocoder.input_dtype == torch.float32
assert vocoder.conv_output_dtype == torch.float32
assert output.dtype == torch.bfloat16
@@ -38,6 +38,7 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
pipeline.default_audio_shift = 3.0 pipeline.default_audio_shift = 3.0
pipeline.device = torch.device("cpu") pipeline.device = torch.device("cpu")
pipeline.od_config = SimpleNamespace() pipeline.od_config = SimpleNamespace()
pipeline.text_encoder = object()
cache_spec = object() cache_spec = object()
quality_plan = SimpleNamespace(cache_dit=cache_spec) quality_plan = SimpleNamespace(cache_dit=cache_spec)
pipeline._quality_policy = Mock() pipeline._quality_policy = Mock()
@@ -110,6 +111,40 @@ def test_pipeline_import_registry_and_component_discovery():
assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"] assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
def test_encoder_free_stage_skips_text_encoder_during_dlo_discovery():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.transformer = torch.nn.Linear(1, 1)
pipeline.video_vae = torch.nn.Linear(1, 1)
pipeline.audio_vae = torch.nn.Linear(1, 1)
pipeline.text_encoder = None
pipeline._dit_modules = ["transformer"]
pipeline._encoder_modules = []
pipeline._vae_modules = ["video_vae", "audio_vae"]
discovered = ModuleDiscovery.discover(pipeline)
assert discovered.dit_names == ["transformer"]
assert discovered.encoder_names == []
assert discovered.vae_names == ["video_vae", "audio_vae"]
def test_encoder_free_stage_skips_missing_component_during_dlo_registration():
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
cache = Mock()
video_vae = Mock()
audio_vae = Mock()
pipeline_module._register_dlo_component_cache(cache, None, video_vae, audio_vae)
video_vae.set_omni_component_cache.assert_called_once_with(cache)
audio_vae.set_omni_component_cache.assert_called_once_with(cache)
def _write_partition_index(path, *, partition, tasks): def _write_partition_index(path, *, partition, tasks):
path.mkdir(parents=True) path.mkdir(parents=True)
(path / "model_index.json").write_text( (path / "model_index.json").write_text(
@@ -575,6 +610,7 @@ def _distilled_pipeline(diffuse_calls, base_schedule_by_partition):
pipeline.default_audio_shift = 3.0 pipeline.default_audio_shift = 3.0
pipeline.device = torch.device("cpu") pipeline.device = torch.device("cpu")
pipeline.od_config = SimpleNamespace() pipeline.od_config = SimpleNamespace()
pipeline.text_encoder = Mock()
pipeline._base_schedule_by_partition = schedules pipeline._base_schedule_by_partition = schedules
pipeline._quality_policy = Mock() pipeline._quality_policy = Mock()
pipeline._quality_policy.resolve.return_value = SimpleNamespace(cache_dit=None) pipeline._quality_policy.resolve.return_value = SimpleNamespace(cache_dit=None)
@@ -964,6 +1000,80 @@ def test_minimax_h3_advertises_the_official_ref2va_image_limit():
assert get_diffusion_model_metadata("MiniMaxH3Pipeline").max_multimodal_image_inputs == 9 assert get_diffusion_model_metadata("MiniMaxH3Pipeline").max_multimodal_image_inputs == 9
def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
class FakeEncoderGroup:
rank_in_group = 0
world_size = 2
def all_reduce(self, tensor):
del tensor
config = SimpleNamespace(
hidden_size=8,
num_attention_heads=4,
num_key_value_heads=2,
head_dim=2,
rms_norm_eps=1e-6,
)
with (
patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
):
attention = encoder_module.MiniMaxH3Qwen3VLTextAttention(FakeEncoderGroup(), config, torch.float32)
# This CPU contract test only needs projection shapes and head reshaping.
# Avoid dispatching platform-specific norm and linear kernels (for example,
# rocm_unquantized_gemm when the test suite is collected on ROCm).
attention.q_norm = nn.Identity()
attention.k_norm = nn.Identity()
with torch.no_grad():
attention.qkv_proj.weight.zero_()
attention.o_proj = nn.Linear(
attention.qkv_proj.local_num_heads * attention.head_dim,
config.hidden_size,
bias=False,
)
attn_call: dict[str, object] = {}
def fake_attention(query, key, value):
attn_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape)
return query
monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", fake_attention)
hidden_states = torch.randn(1, 3, config.hidden_size)
cos = torch.ones(1, 3, config.head_dim)
sin = torch.zeros_like(cos)
output = attention(hidden_states, (cos, sin))
assert attn_call["query_shape"] == (1, config.num_attention_heads // 2, 3, config.head_dim)
assert attn_call["key_shape"] == (1, config.num_key_value_heads // 2, 3, config.head_dim)
assert attn_call["value_shape"] == (1, config.num_key_value_heads // 2, 3, config.head_dim)
assert output.shape == hidden_states.shape
def test_text_attention_sdpa_helper_preserves_expanded_kv_fallback(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
sdpa_call: dict[str, object] = {}
def fake_sdpa(query, key, value, **kwargs):
sdpa_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape, kwargs=kwargs)
return query
monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)
output = encoder_module._scaled_dot_product_attention(
torch.randn(1, 4, 3, 8),
torch.randn(1, 2, 3, 8),
torch.randn(1, 2, 3, 8),
)
assert sdpa_call["query_shape"] == sdpa_call["key_shape"] == sdpa_call["value_shape"] == (1, 4, 3, 8)
assert sdpa_call["kwargs"] == {"dropout_p": 0.0, "is_causal": True}
assert output.shape == (1, 4, 3, 8)
def test_encoder_forward_uses_hook_compatible_encode_entrypoint(): def test_encoder_forward_uses_hook_compatible_encode_entrypoint():
from vllm_omni.diffusion.models.minimax_h3.encoder import ( from vllm_omni.diffusion.models.minimax_h3.encoder import (
MiniMaxH3Qwen3VLEncoder, MiniMaxH3Qwen3VLEncoder,
@@ -1021,7 +1131,7 @@ def test_encoder_forward_forwards_video_inputs():
def test_reference_video_shape_uses_h3_adapt_shape_policy(): def test_reference_video_shape_uses_h3_adapt_shape_policy():
from vllm_omni.diffusion.models.minimax_h3.reference_video import ( from vllm_omni.model_executor.models.minimax_h3.reference_video import (
_reference_video_shape, _reference_video_shape,
) )
@@ -1607,6 +1717,85 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti
assert transcode_calls[0][1]["target_frame_count"] == 209 assert transcode_calls[0][1]["target_frame_count"] == 209
def test_ref2va_transcode_zero_frame_count_keeps_video_stream(monkeypatch, tmp_path):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
commands: list[list[str]] = []
monkeypatch.setattr(
reference_video_module.subprocess,
"run",
lambda command, **_kwargs: commands.append(command),
)
reference_video_module._transcode_reference_video(
"input.mp4",
target_width=768,
target_height=448,
target_frame_count=0,
workdir=str(tmp_path),
duration_seconds=4.0,
)
assert "-frames:v" not in commands[0]
def test_prepared_reference_video_descriptor_round_trip(tmp_path):
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
deserialize_prepared_reference_videos,
serialize_prepared_reference_videos,
)
item = {
"original_path": str(tmp_path / "source.mp4"),
"prepared_path": str(tmp_path / "prepared.mp4"),
"input_has_audio": True,
"width": 1344,
"height": 768,
"start_time_seconds": 0.0,
"duration_seconds": 5.2,
"audio_duration_seconds": 5.2,
}
descriptor = serialize_prepared_reference_videos([item], str(tmp_path))
assert deserialize_prepared_reference_videos(descriptor) == (str(tmp_path), [item])
def test_prepared_reference_video_descriptor_rejects_unknown_fields():
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
deserialize_prepared_reference_videos,
)
with pytest.raises(ValueError, match="invalid MiniMax H3 prepared-reference-video descriptor"):
deserialize_prepared_reference_videos('{"artifact_dir":"/tmp","videos":[{"path":"x"}]}')
def test_stage_one_reuses_existing_prepared_reference_video(tmp_path):
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_reuse_prepared_reference_videos,
)
prepared_path = tmp_path / "prepared.mp4"
prepared_path.touch()
prepared = [{"prepared_path": str(prepared_path)}]
assert _reuse_prepared_reference_videos(prepared, expected_count=1) is prepared
def test_stage_one_rejects_invalid_prepared_reference_video(tmp_path):
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_reuse_prepared_reference_videos,
)
with pytest.raises(ValueError, match="count does not match"):
_reuse_prepared_reference_videos([], expected_count=1)
with pytest.raises(ValueError, match="prepared reference video is unavailable"):
_reuse_prepared_reference_videos(
[{"prepared_path": str(tmp_path / "missing.mp4")}],
expected_count=1,
)
def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch): def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
@@ -2046,7 +2235,7 @@ def test_g4_reference_image_file_format_and_size_contract(tmp_path):
def test_g4_standalone_audio_duration_and_total_duration_contract(): def test_g4_standalone_audio_duration_and_total_duration_contract():
from vllm_omni.diffusion.models.minimax_h3.reference_video import ( from vllm_omni.model_executor.models.minimax_h3.reference_video import (
validate_reference_audio_waveforms, validate_reference_audio_waveforms,
) )
@@ -2121,7 +2310,7 @@ def test_ref2va_standalone_audio_condition_is_bounded_to_output_duration():
], ],
) )
def test_g4_reference_video_metadata_validation(field, value, message, tmp_path): def test_g4_reference_video_metadata_validation(field, value, message, tmp_path):
from vllm_omni.diffusion.models.minimax_h3.reference_video import ( from vllm_omni.model_executor.models.minimax_h3.reference_video import (
_validate_reference_video_metadata, _validate_reference_video_metadata,
) )
@@ -197,7 +197,6 @@ def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path):
[ [
"model-level CPU offload (--enable-cpu-offload)", "model-level CPU offload (--enable-cpu-offload)",
"layerwise offload (--enable-layerwise-offload)", "layerwise offload (--enable-layerwise-offload)",
"distributed layerwise offload (--enable-distributed-layerwise-offload)",
], ],
) )
def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode): def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
@@ -214,6 +213,38 @@ def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
) )
def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.models.minimax_h3 import (
pipeline_minimax_h3 as pipeline_module,
)
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline._turbo_lora_adapter_ids = set()
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
enable_distributed_layerwise_offload=True,
)
captured = {}
def load_turbo(**kwargs):
captured.update(kwargs)
return object(), object()
monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo)
loaded = pipeline._load_diffusion_lora_adapter(
lora_request=_request("turbo"),
lora_path="turbo",
dtype=torch.bfloat16,
)
assert loaded is not None
assert captured["unsupported_offload_mode"] is None
def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path): def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path):
unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors" unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors"
_write_tiny_turbo(unsupported) _write_tiny_turbo(unsupported)
@@ -254,18 +254,3 @@ def test_pipeline_resolves_transformer_component_quant_config():
assert transformer_config.ignored_layers == ignored_layers assert transformer_config.ignored_layers == ignored_layers
assert _resolve_component_quant_config(component_config, "transformer") is transformer_config assert _resolve_component_quant_config(component_config, "transformer") is transformer_config
assert _resolve_component_quant_config(transformer_config, "transformer") is transformer_config assert _resolve_component_quant_config(transformer_config, "transformer") is transformer_config
def test_pipeline_strips_prequantized_text_encoder_config():
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _resolve_minimax_h3_text_encoder_quant_config
from vllm_omni.quantization import ComponentQuantizationConfig
prequantized = Mock()
prequantized.get_name.return_value = "modelopt"
online_fp8 = Mock()
online_fp8.get_name.return_value = "fp8"
component_config = ComponentQuantizationConfig({"text_encoder": prequantized})
assert _resolve_minimax_h3_text_encoder_quant_config(online_fp8) is online_fp8
assert _resolve_minimax_h3_text_encoder_quant_config(prequantized) is None
assert _resolve_minimax_h3_text_encoder_quant_config(component_config) is None
@@ -187,7 +187,6 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
"model": model_ref, "model": model_ref,
"enforce_eager": True, "enforce_eager": True,
"tensor_parallel_size": 2, "tensor_parallel_size": 2,
"text_encoder_tp_size": 2,
"vae_use_tiling": True, "vae_use_tiling": True,
} }
@@ -0,0 +1,132 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import torch.nn.functional as F
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
def test_minimax_h3_npu_sdpa_patch_is_independent_and_idempotent(monkeypatch) -> None:
from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
def sdpa_sentinel(query, key, value):
del key, value
return query
def rope_sentinel(q, k, cos, sin):
del cos, sin
return q, k
monkeypatch.setattr(npu_minimax_h3, "_ROPE_PATCHED", False)
monkeypatch.setattr(npu_minimax_h3, "_SDPA_PATCHED", False)
monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel)
monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel)
npu_minimax_h3.apply_minimax_h3_qwen3vl_sdpa_patch()
assert encoder_module._scaled_dot_product_attention is npu_minimax_h3._scaled_dot_product_attention_npu
assert encoder_module._apply_rotary_pos_emb is rope_sentinel
assert not npu_minimax_h3._ROPE_PATCHED
assert npu_minimax_h3._SDPA_PATCHED
monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel)
npu_minimax_h3.apply_minimax_h3_qwen3vl_sdpa_patch()
assert encoder_module._scaled_dot_product_attention is sdpa_sentinel
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
@pytest.mark.parametrize(
("num_heads", "num_kv_heads", "expected_enable_gqa"),
[(2, 2, False), (4, 2, True)],
)
def test_minimax_h3_npu_sdpa_preserves_compressed_heads_and_selects_gqa(
monkeypatch,
num_heads: int,
num_kv_heads: int,
expected_enable_gqa: bool,
) -> None:
from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
captured: dict[str, object] = {}
def fake_sdpa(query, key, value, **kwargs):
captured.update(query=query, key=key, value=value, kwargs=kwargs)
return query
monkeypatch.setattr(F, "scaled_dot_product_attention", fake_sdpa)
query = torch.randn(1, num_heads, 3, 8)
key = torch.randn(1, num_kv_heads, 3, 8)
value = torch.randn_like(key)
output = npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value)
assert captured["query"] is query
assert captured["key"] is key
assert captured["value"] is value
assert captured["kwargs"] == {
"dropout_p": 0.0,
"is_causal": True,
"enable_gqa": expected_enable_gqa,
}
assert output is query
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
@pytest.mark.parametrize(
("num_heads", "num_key_heads", "num_value_heads", "message"),
[
(4, 2, 1, "key and value to have the same number of heads"),
(4, 0, 0, "at least one KV head"),
(3, 2, 2, "query heads to be a multiple of KV heads"),
],
)
def test_minimax_h3_npu_sdpa_rejects_invalid_head_configurations(
num_heads: int,
num_key_heads: int,
num_value_heads: int,
message: str,
) -> None:
from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
query = torch.randn(1, num_heads, 3, 8)
key = torch.randn(1, num_key_heads, 3, 8)
value = torch.randn(1, num_value_heads, 3, 8)
with pytest.raises(ValueError, match=message):
npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value)
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
def test_minimax_h3_npu_causal_gqa_matches_expanded_kv() -> None:
"""Exercise the real torch-npu causal GQA operator."""
from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
torch.manual_seed(0)
batch_size, seq_len = 1, 16
num_heads, num_kv_heads, head_size = 8, 2, 128
device = torch.device("npu")
dtype = torch.bfloat16
query = torch.randn(batch_size, num_heads, seq_len, head_size, device=device, dtype=dtype)
key = torch.randn(batch_size, num_kv_heads, seq_len, head_size, device=device, dtype=dtype)
value = torch.randn_like(key)
actual = npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value)
repeat_num = num_heads // num_kv_heads
expected = F.scaled_dot_product_attention(
query,
key.repeat_interleave(repeat_num, dim=1),
value.repeat_interleave(repeat_num, dim=1),
dropout_p=0.0,
is_causal=True,
)
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
@@ -6,7 +6,38 @@ import torch
from vllm_omni.platforms import current_omni_platform from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
def test_qwen3_vl_npu_rope_patch_is_independent_and_idempotent(monkeypatch) -> None:
from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
def rope_sentinel(q, k, cos, sin):
del cos, sin
return q, k
def sdpa_sentinel(query, key, value):
del key, value
return query
monkeypatch.setattr(npu_minimax_h3, "_ROPE_PATCHED", False)
monkeypatch.setattr(npu_minimax_h3, "_SDPA_PATCHED", False)
monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel)
monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel)
npu_minimax_h3.apply_minimax_h3_qwen3vl_patch()
assert encoder_module._apply_rotary_pos_emb is npu_minimax_h3._apply_rotary_pos_emb_npu
assert encoder_module._scaled_dot_product_attention is sdpa_sentinel
assert npu_minimax_h3._ROPE_PATCHED
assert not npu_minimax_h3._SDPA_PATCHED
monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel)
npu_minimax_h3.apply_minimax_h3_qwen3vl_patch()
assert encoder_module._apply_rotary_pos_emb is rope_sentinel
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") @pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
@@ -1,13 +1,18 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest import pytest
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
from tests.helpers.mark import hardware_marks
from vllm_omni.platforms import current_omni_platform from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] pytestmark = [
pytest.mark.core_model,
pytest.mark.diffusion,
*hardware_marks(res={"npu": "A3"}, num_cards=1),
]
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") @pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
Loaded 100 of 381 files, more files were not shown because too many files have changed in this diff. Show more