Compare commits

..
Author SHA1 Message Date
Yueqian Lin aef61597d3 Merge remote-tracking branch 'origin/main' into HEAD 2026-08-30 03:36:08 -07:00
akshatvishu c41925689d [BugFix][Higgs-Audio-V3] Fix transcript and concurrent sampling tests (#6422)
Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
2026-08-30 08:52:15 +00:00
Yueqian Lin a6b559cf54 [Bugfix][MiniMax-Music3] Resolve stage subdirs against a real snapshot (#6640)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-08-30 06:51:24 +00:00
baonudesifeizhai 8be601ed7f [Diffusion] Add native SymmMem Fast Ulysses transport (#6340)
Signed-off-by: roG0d <baonudesifeizhai@gmail.com>
2026-08-30 04:14:21 +00:00
y-null 99a2ae1fc8 [Bugfix][MiniCPM-o] Anchor duplex soft-interrupt listen sandwich before commit (#6772)
Signed-off-by: y_null <y_null@qq.com>
2026-08-30 03:25:43 +00:00
21dd4b33ac [Bugfix][MiniCPM-o] Stabilize native duplex streaming and Stage-1 handoffs (#6529)
Signed-off-by: Sy03 <1370724210@qq.com>
Co-authored-by: Nathan Price <nathan@abridge.com>
Co-authored-by: Eric(Qun) <186020953+0xtoward@users.noreply.github.com>
2026-08-29 13:05:24 -07:00
WeiQing ChenandGao Han 9a88d2d21d [diffusion][bugfix] Separate Diffusers hooks from model metadata (#6749)
Signed-off-by: david6666666 <530634352@qq.com>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-08-29 17:37:14 +00:00
YancyandAsthenia 73893e73ea [Bugfix] Fix create_error_response import across vLLM versions (#6773)
Signed-off-by: Asthenia <asthenia0412@gmail.com>
Co-authored-by: Asthenia <asthenia0412@gmail.com>
2026-08-29 17:01:51 +00:00
NATURE 87736a3cc5 [Bugfix][MiniCPM-o] Bound full-duplex Talker context with sliding recompute (#6626)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
2026-08-30 00:49:42 +08:00
Sy03 3adee400a0 [Bugfix][Qwen3-TTS] Handle codec generations that exhaust their token budget (#6728)
Signed-off-by: Sy03 <1370724210@qq.com>
2026-08-29 08:59:05 -07:00
Gao Han 676a35fa8b Fix MiniCPM-o concurrent audio first-packet latency (#6767)
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
2026-08-29 23:27:47 +08:00
Honghan Zhuandamy-why-3459 72a1ce488b [Bugfix] Avoid invalid TPOT metrics (#6696)
Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-29 11:56:37 +00:00
wangyu e338f0ece3 [CI]Split GPU jobs by cards_* and reject hand-written SKU marks (#6650)
Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: [Your Name] <your.email@example.com>
2026-08-29 18:42:33 +08:00
HAAZZZEEEE c50b7de167 [NPU][Diffusion] MiniMax-H3 RainFusion end_step tail fallback (#6037)
Signed-off-by: HAAZZZEEEE <wang-taicheng@qq.com>
2026-08-29 18:36:05 +08:00
psv666 a5b8c82788 [Bugfix][MiniCPM-o] Fix duplex camera frame fixture (#6757)
Signed-off-by: psv666 <2693925048@qq.com>
2026-08-29 09:57:53 +00:00
ooooooye 5a5bc55c25 [Feature][Diffusion] Allow online INT8 quantization with DLO AllGather (#6573)
Signed-off-by: brandneway <gyuan4892@gmail.com>
2026-08-29 09:50:30 +00:00
Wenjia LiandWeiQing Chen d9980c75e8 Perf/minimax h3 dit swiglu rope cache npu (#6410)
Signed-off-by: Wenjia Li <wjialish@gmail.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
2026-08-29 09:14:52 +00:00
bb5bf8fd0a [Bugfix] Fix colocate-async sleep admission race and deliver AR abort tokens (#6367)
Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Signed-off-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Samit <285365963@qq.com>
2026-08-29 09:11:39 +00:00
psv666andGao Han 5c47fd0bc8 [Bugfix][Qwen3-Omni] Handle missing packed modules mapping (#6748)
Signed-off-by: psv666 <2693925048@qq.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-08-29 15:24:49 +08:00
Zijun HuangandWeiQing Chen 08cdee8dc4 [Feature] Support MiniMax-H3 FlashGen native LoRA with the legacy manager (#6666)
Signed-off-by: Zijun Huang <zijunhua@usc.edu>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
2026-08-29 06:52:19 +00:00
143ccc5a2f Fix Qwen3-Omni audio encoder TP when heads are not divisible by TP size (#4322)
Signed-off-by: abrahamzewoudie <azewoudi@amd.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-29 14:52:02 +08:00
WeiQing Chen 55a226dcf1 [video][bugfix] Accept semantic video output in offline examples (#6747)
Signed-off-by: david6666666 <530634352@qq.com>
2026-08-29 06:26:28 +00:00
Jiatai Wang f66f05633c [Bugfix] Honor request seed in MOSS-TTS adapters (#6543)
Signed-off-by: JiataiWang <wangjiatai@proton.me>
2026-08-28 22:44:13 -07:00
ZacheryAUandamy-why-3459 457c6e2793 [bugfix][MiniCPM-o] Re-land empty full-payload and async_chunk=false processor tests (#6745)
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-29 13:39:54 +08:00
boatman bfe8967ca2 [Bugfix][StepAudio2] Fix async-chunk metadata mismatch and boundary glitches (#5917)
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
2026-08-28 21:45:24 -07:00
WeiQing ChenandOpenAI Codex fbc93a0e08 [ci][bugfix] Fix MiniMax-H3 FP8 quality test OOM (#6742)
Signed-off-by: david6666666 <530634352@qq.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-08-29 12:11:31 +08:00
Bob Tian d934a5f7b8 [Bugfix] Report engine-queued requests as waiting in vllm_omni:num_requests gauges (#6549)
Signed-off-by: Bob Tian <bobzetian@google.com>
2026-08-29 03:58:26 +00:00
Inesh Reddy Chappidiandamy-why-3459 bbb3ae01a1 [CI] Gate qwen3-omni no-async-chunk perf on mean_audio_rtf (#6743)
Signed-off-by: Inesh Reddy <ineshreddy249@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-28 20:28:58 -07:00
Zhang BowenandHongsheng Liu 46cfa8435b [Model][Frontend] MiniMax-H3: Parallelize MP4 response frame conversion (#6499)
Signed-off-by: MosCloud <bwzhang1991@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-28 21:22:00 -06:00
fan2956andWeiQing Chen 0dc8f45906 [Feat] Minimax H3 support ref and tgt video sparse attention on NPU (#6518)
Signed-off-by: fan2956 <zhoufan53@huawei.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
2026-08-29 11:10:53 +08:00
d94061201d [Bugfix][Engine] Make abort idempotent during shutdown (#6327)
Signed-off-by: Allen Wu <allenwu2795@gmail.com>
Co-authored-by: TRAE CLI <noreply@bytedance.com>
Co-authored-by: TRAE CLI <traecli@bytedance.com>
2026-08-28 19:51:25 -07:00
Mu GuanLinandWeiQing Chen 9558fd2ea8 [Bugfix] Scope MiniMax-H3 cuDNN SDPA state (#6710)
Signed-off-by: mglyn <1203789601@qq.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
2026-08-28 19:40:50 -07:00
Mu GuanLin 55a8bd6b80 [CI][Bugfix] Align the LTX-2 FP8 quality gate with the default recipe (#5831)
Signed-off-by: mglyn <1203789601@qq.com>
2026-08-29 02:24:09 +00:00
Honghan Zhu f584729f24 [BugFix] Fix local media access for diffusion speech (#6622)
Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>
2026-08-29 02:01:24 +00:00
Chendi.Xue 034af1f657 [CI][XPU]switch to use vLLM base docker for XPU (#6727)
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
2026-08-29 09:03:00 +08:00
Hongsheng Liuandhsliuustc0106 8bfad23f98 [Core] Restore topology-aware diffusion worker titles (#6722)
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
2026-08-29 08:45:57 +08:00
Gao Han 6b9ce05feb Revert "[bugfix][MiniCPM-o] Fix offline_inference/test_minicpmo_4_5.py and online_serving/ one" (#6730) 2026-08-29 08:14:02 +08:00
ZacheryAUandamy-why-3459 6d84409a68 [bugfix][MiniCPM-o] Fix offline_inference/test_minicpmo_4_5.py and online_serving/ one (#5464)
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-28 19:20:09 +00:00
ebb3b20576 [Diffusion] Make Flux2Pipeline text_encoder_out_layers configurable (#6390)
Co-authored-by: khairulkabir1661 <khairulkabir1661@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-08-28 09:05:41 -06:00
be335a86f9 fix(npu): extend bf16 autocast to NPU in MossAudioTokenizer decode (#6664)
Signed-off-by: Wallbreazzz <110282866+Wallbreazzz@users.noreply.github.com>
Signed-off-by: Canlin Guo <961750412@qq.com>
Co-authored-by: Canlin Guo <961750412@qq.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>
2026-08-28 04:01:49 -07:00
Hongsheng Liu d5aced9296 Revert "[Core] Make diffusion worker titles topology-aware" (#6717) 2026-08-28 17:30:03 +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
Yueqian Lin 1ad90653c5 [CI/Build] Add dots.tts weekly e2e coverage
dots.tts landed in #4765 with no GPU coverage: nothing under tests/e2e
exercises it, so the DiT sampler, streaming vocoder and AR loopback have
no CI signal at all.

Add an offline e2e module marked slow + tts + L4, which the existing
weekly 'TTS · L4' step collects via -m 'slow and L4 and tts'. No pipeline
YAML change is needed; the weekly TTS steps are marker driven.

Three cases, the latter two guarding specific review findings from #4765:
  1. zero-shot synthesis produces finite, non-silent audio
  2. mixed prefill+decode batch keeps per-request logits rows aligned
  3. the same text twice yields the same audio (seeded FM noise)

Verified selector behaviour: 3 collected under the weekly selector,
0 collected under the nightly one.

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-08-13 11:24:50 -07:00
344 changed files with 25571 additions and 2306 deletions

No files matched your search

+89 -20
View File
@@ -1,4 +1,7 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# vllm-omni customized version
# Based on: https://github.com/vllm-project/ci-infra/blob/main/buildkite/bootstrap-amd.sh
# Last synced: 2025-12-15
@@ -6,6 +9,9 @@
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
if [[ -z "${RUN_ALL:-}" ]]; then
@@ -63,6 +69,8 @@ upload_pipeline() {
# Install minijinja
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
# Installed by the minijinja bootstrap above and only present on the CI agent.
# shellcheck disable=SC1091
source /var/lib/buildkite-agent/.cargo/env
if [[ $BUILDKITE_PIPELINE_SLUG == "fastcheck" ]]; then
@@ -90,22 +98,82 @@ upload_pipeline() {
cd .buildkite/amd
# Select test definition file: merge suite for main, ready suite for PRs.
# For debugging, DEBUG_TEST_YAML can override the selection — accepts
# "merge" or "ready" (case-insensitive).
# For debugging, DEBUG_TEST_YAML accepts a comma-separated list containing
# "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
case "${DEBUG_TEST_YAML,,}" in
merge)
TEST_YAML="test-amd-merge.yml"
;;
ready)
TEST_YAML="test-amd-ready.yml"
;;
*)
echo "ERROR: DEBUG_TEST_YAML must be 'merge' or 'ready', got '$DEBUG_TEST_YAML'" >&2
declare -a DEBUG_TEST_SPECS=()
declare -A SEEN_DEBUG_TESTS=()
IFS=',' read -ra REQUESTED_DEBUG_TESTS <<< "${DEBUG_TEST_YAML,,}"
for requested_test in "${REQUESTED_DEBUG_TESTS[@]}"; do
# Permit readable values such as "ready, merge," and ignore the
# empty item produced by a trailing comma.
requested_test="${requested_test//[[:space:]]/}"
[[ -z "$requested_test" ]] && continue
if [[ -n "${SEEN_DEBUG_TESTS[$requested_test]:-}" ]]; then
echo "ERROR: duplicate DEBUG_TEST_YAML suite '$requested_test'" >&2
exit 1
;;
esac
echo "DEBUG_TEST_YAML override: using $TEST_YAML"
fi
SEEN_DEBUG_TESTS[$requested_test]=1
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
TEST_YAML="test-amd-merge.yml"
else
@@ -130,19 +198,20 @@ upload_pipeline() {
> 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 pipeline upload pipeline.yaml
exit 0
}
get_diff() {
$(git add .)
echo $(git diff --name-only --diff-filter=ACMDR $(git merge-base origin/main HEAD))
git diff --name-only --diff-filter=ACMDR "$(git merge-base origin/main HEAD)"
}
get_diff_main() {
$(git add .)
echo $(git diff --name-only --diff-filter=ACMDR HEAD~1)
git diff --name-only --diff-filter=ACMDR HEAD~1
}
file_diff=$(get_diff)
@@ -181,7 +250,7 @@ for file in $file_diff; do
# First check if file matches any pattern
matches_pattern=0
for pattern in "${patterns[@]}"; do
if [[ $file == $pattern* ]] || [[ $file == $pattern ]]; then
if [[ $file == "$pattern"* ]] || [[ $file == "$pattern" ]]; then
matches_pattern=1
break
fi
@@ -191,7 +260,7 @@ for file in $file_diff; do
if [[ $matches_pattern -eq 1 ]]; then
matches_ignore=0
for ignore in "${ignore_patterns[@]}"; do
if [[ $file == $ignore* ]] || [[ $file == $ignore ]]; then
if [[ $file == "$ignore"* ]] || [[ $file == "$ignore" ]]; then
matches_ignore=1
break
fi
+45 -17
View File
@@ -21,7 +21,7 @@ steps:
- label: "Simple · Diffusion Test · Shard %N/%t"
agent_pool: mi300_1
parallelism: 2
parallelism: 4
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -71,26 +71,30 @@ steps:
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
bash -c '
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))
'
- pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
- pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs) · Shard 1/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
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
- timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
- pytest -s -v tests/diffusion/distributed/test_sequence_parallel.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
- label: "Diffusion Tensor Parallelism Test"
@@ -171,21 +175,45 @@ steps:
grade: Blocking
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model'
- label: "Qwen3-TTS CustomVoice E2E Test · Shard %N/%t"
agent_pool: mi300_1
parallelism: 2
parallelism: 3
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
timeout_in_minutes: 80
commands:
- |
bash -c '
timeout --signal=TERM --kill-after=2m 80m 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 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"
+1 -4
View File
@@ -191,10 +191,7 @@ steps:
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 45m bash -c '
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model"
'
- 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
# Need to debug before reneable numerical changes across large PRs
+16 -1
View File
@@ -122,6 +122,7 @@ steps:
- label: "Omni · Qwen3-Omni Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_qwen3_omni.py
- tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py
- tests/e2e/online_serving/test_qwen3_omni.py
- tests/helpers/runtime.py
- vllm_omni/entrypoints/openai/serving_chat.py
@@ -132,9 +133,10 @@ steps:
- vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
- vllm_omni/deploy/qwen3_omni_moe.yaml
- vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
timeout_in_minutes: 50
commands:
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: h100_2
- label: "Omni · MiniCPM-o 4.5 Test"
@@ -166,6 +168,7 @@ steps:
- tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.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_server_vad.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
- vllm_omni/deploy/minicpmo_4_5.yaml
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
@@ -216,3 +219,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'
"
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
+131 -34
View File
@@ -12,16 +12,22 @@ steps:
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test"
steps:
- label: ":full_moon: Omni · Function Test with H100"
timeout_in_minutes: 120
- label: ":full_moon: Omni · Function Test with H100 · Single-GPU"
timeout_in_minutes: 90
commands:
- pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py --ignore=tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
- pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni and cards_1" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
mirror_hardwares: h100_1
- label: ":full_moon: Omni · Function Test with H100 · 2-GPU"
timeout_in_minutes: 90
commands:
- pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni and cards_2" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py
mirror_hardwares: h100_2
- label: ":full_moon: Omni · MiniCPM-o 4.5 Duplex Test"
timeout_in_minutes: 50
commands:
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni" --run-level "full_model"
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni and cards_1" --run-level "full_model"
mirror_hardwares: h100_1
- label: ":full_moon: Omni · Doc Test with L4"
@@ -31,10 +37,10 @@ steps:
- pytest -s -v tests/examples/ -m "full_model and omni and L4" --run-level "full_model"
mirror_hardwares: l4_4
- label: ":full_moon: Omni · Doc Test with H100"
- label: ":full_moon: Omni · Doc Test with H100 · 2-GPU"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/examples/ -m "full_model and omni and H100" --run-level "full_model"
- pytest -s -v tests/examples/ -m "full_model and omni and H100 and cards_2" --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Omni · Accuracy Test"
@@ -136,15 +142,29 @@ steps:
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
mirror_hardwares: l4_1
- label: ":full_moon: TTS · Perf Test"
key: nightly-tts-performance
- label: ":full_moon: TTS · Perf Test · Single-GPU"
key: nightly-tts-performance-single
timeout_in_minutes: 180
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json -m "H100 and cards_1"
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
mirror_hardwares: h100_1
- label: ":full_moon: TTS · Perf Test · 2-GPU"
key: nightly-tts-performance-2gpu
timeout_in_minutes: 180
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json -m "H100 and cards_2"
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
@@ -179,7 +199,6 @@ steps:
build.pull_request.labels includes "nightly-test" ||
build.pull_request.labels includes "diffusion-x2iat-test"
steps:
# Use explicit file shards so single-card cases do not reserve 2+ GPUs.
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Single-GPU"
timeout_in_minutes: 120
@@ -187,24 +206,24 @@ steps:
- >-
pytest -sv
tests/e2e/online_serving/test_qwen_image_expansion.py
-m "full_model and diffusion and H100 and not distributed_cuda"
-m "full_model and diffusion and H100 and cards_1"
--run-level "full_model"
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Qwen-Image"
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · 2-GPU"
timeout_in_minutes: 120
commands:
- >-
pytest -sv
tests/e2e/online_serving/test_qwen_image_expansion.py
-m "full_model and diffusion and H100 and distributed_cuda"
-m "full_model and diffusion and H100 and cards_2"
--run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4 · Multi-GPU"
timeout_in_minutes: 120
commands:
- pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
- pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4 and not cards_1" --run-level "full_model" --ignore=tests/e2e/accuracy
mirror_hardwares: l4_4
- label: ":full_moon: Diffusion X2I(&A&T) · Doc Test"
@@ -235,8 +254,8 @@ steps:
- pytest -s -v tests/e2e/accuracy/test_hidream_o1_image.py -m full_model --run-level full_model
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image"
key: nightly-diffusion-x2iat-performance-qwen-image
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · Single-GPU"
key: nightly-diffusion-x2iat-performance-qwen-image-single
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
@@ -244,7 +263,23 @@ steps:
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json -m "H100 and cards_1"
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · 4-GPU"
key: nightly-diffusion-x2iat-performance-qwen-image-4gpu
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json -m "H100 and cards_4"
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
@@ -260,16 +295,34 @@ steps:
build.pull_request.labels includes "diffusion-x2v-test"
steps:
- label: ":full_moon: Diffusion X2V · Wan2.2 T2V Function Test"
- label: ":full_moon: Diffusion X2V · Wan2.2 T2V Function Test · Single-GPU"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda" --run-level "full_model"
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda and H100 and cards_1" --run-level "full_model"
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2V · Wan2.2 T2V Function Test · 2-GPU"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda and H100 and cards_2" --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test"
- label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test · Single-GPU"
timeout_in_minutes: 90
commands:
- 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 and H100 and cards_1" --run-level "full_model"
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test · 2-GPU"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda and H100 and cards_2" --run-level "full_model"
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"
@@ -304,17 +357,49 @@ steps:
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · Perf Test"
key: nightly-diffusion-x2v-performance
- label: ":full_moon: Diffusion X2V · Perf Test · Single-GPU"
key: nightly-diffusion-x2v-performance-single
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json -m "H100 and cards_1"
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2V · Perf Test · 2-GPU"
key: nightly-diffusion-x2v-performance-2gpu
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json -m "H100 and cards_2"
EXIT1=$$?
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json -m "H100 and cards_2"
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 · 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"
@@ -348,10 +433,16 @@ steps:
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4' --run-level "full_model"
mirror_hardwares: l4_4
- label: ":full_moon: Diffusion · Distributed Test with H100"
- label: ":full_moon: Diffusion · Distributed Test with H100 · 2-GPU"
timeout_in_minutes: 60
commands:
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100' --run-level "full_model"
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100 and cards_2' --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion · Distributed Test with H100 · 4-GPU"
timeout_in_minutes: 60
commands:
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100 and cards_4' --run-level "full_model"
mirror_hardwares: h100_4
- label: ":full_moon: Diffusion · Offloader Test"
@@ -386,19 +477,25 @@ steps:
depends_on:
- nightly-omni-performance-no-async-chunk
- nightly-omni-performance-async-chunk
- nightly-tts-performance
- nightly-diffusion-x2iat-performance-qwen-image
- nightly-diffusion-x2v-performance
- nightly-tts-performance-single
- nightly-tts-performance-2gpu
- nightly-diffusion-x2iat-performance-qwen-image-single
- nightly-diffusion-x2iat-performance-qwen-image-4gpu
- nightly-diffusion-x2v-performance-single
- nightly-diffusion-x2v-performance-2gpu
if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
commands:
- pip install openpyxl
- export DEFAULT_INPUT_DIR=tests/dfx/perf/results
- export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-single
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-2gpu
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-single
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-4gpu
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance-single
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance-2gpu
- python tools/nightly/generate_nightly_perf_excel.py
- python tools/nightly/generate_nightly_perf_html.py
- python tools/nightly/send_nightly_email.py --report-file "tests/dfx/perf/results/*.xlsx, tests/dfx/perf/results/*.html"
+13
View File
@@ -205,6 +205,7 @@ steps:
- tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.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_server_vad.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
- vllm_omni/deploy/minicpmo_4_5.yaml
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
@@ -337,3 +338,15 @@ steps:
pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' --run-level 'core_model'
"
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
+48 -14
View File
@@ -44,10 +44,16 @@ steps:
- pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
mirror_hardwares: h100_4
- label: "Reliability Test · Invalid parameters · H100"
timeout_in_minutes: 120
- label: "Reliability Test · Invalid parameters · H100 · Single-GPU"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100"
- pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100 and cards_1"
mirror_hardwares: h100_1
- label: "Reliability Test · Invalid parameters · H100 · 2-GPU"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100 and cards_2"
mirror_hardwares: h100_2
- label: "Reliability Test · Invalid parameters · L4"
@@ -105,27 +111,37 @@ steps:
if: build.env("NON_CRITICAL") == "1"
steps:
- label: "Omni · L4"
# Qwen2.5 expansion (CUDA 4-card / online 2-card) stays on this l4_4 job.
- label: "Omni · L4 · Multi-GPU"
timeout_in_minutes: 90
commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py -m "slow and L4 and omni and not cards_1" --run-level "core_model"
- >-
pytest -s -v tests/e2e/
--ignore=tests/e2e/online_serving/test_qwen2_5_omni_expansion.py
--ignore=tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
-m "slow and L4 and omni"
-m "slow and L4 and omni and not cards_1"
--run-level "full_model"
mirror_hardwares: l4_4
- label: "Omni · H100"
- label: "Omni · H100 · Single-GPU"
timeout_in_minutes: 90
commands:
- >-
pytest -sv
tests/e2e/ -m "slow and H100 and omni"
tests/e2e/ -m "slow and H100 and omni and cards_1"
--run-level "full_model"
mirror_hardwares: h100_2
mirror_hardwares: h100_1
- label: "Omni · H100 · 4-GPU"
timeout_in_minutes: 90
commands:
- >-
pytest -sv
tests/e2e/ -m "slow and H100 and omni and cards_4"
--run-level "full_model"
mirror_hardwares: h100_4
- label: "TTS · L4"
timeout_in_minutes: 120
@@ -169,12 +185,21 @@ steps:
--run-level "full_model"
mirror_hardwares: l4_1
- label: "Diffusion · L4"
- label: "Diffusion · L4 · Single-GPU"
timeout_in_minutes: 60
commands:
- >-
pytest -sv
tests/e2e/ -m "slow and diffusion and L4"
tests/e2e/ -m "slow and diffusion and L4 and cards_1"
--run-level "full_model"
mirror_hardwares: l4_1
- label: "Diffusion · L4 · Multi-GPU"
timeout_in_minutes: 60
commands:
- >-
pytest -sv
tests/e2e/ -m "slow and diffusion and L4 and not cards_1"
--run-level "full_model"
mirror_hardwares: l4_4
@@ -183,15 +208,24 @@ steps:
commands:
- >-
pytest -sv
tests/e2e/ -m "slow and diffusion and H100 and not distributed_cuda"
tests/e2e/ -m "slow and diffusion and H100 and cards_1"
--run-level "full_model"
mirror_hardwares: h100_1
- label: "Diffusion · H100 · Multi-GPU"
- label: "Diffusion · H100 · 2-GPU"
timeout_in_minutes: 90
commands:
- >-
pytest -sv
tests/e2e/ -m "slow and diffusion and H100 and distributed_cuda"
tests/e2e/ -m "slow and diffusion and H100 and cards_2"
--run-level "full_model"
mirror_hardwares: h100_2
- label: "Diffusion · H100 · 4-GPU"
timeout_in_minutes: 90
commands:
- >-
pytest -sv
tests/e2e/ -m "slow and diffusion and H100 and cards_4"
--run-level "full_model"
mirror_hardwares: h100_4
+1 -1
View File
@@ -10,7 +10,7 @@ steps:
DOCKER_BUILDKIT: "1"
# Buildkite will automatically replace this with the actual commit hash
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
VLLM_VERSION: "v0.27.0"
VLLM_VERSION: "v0.28.0"
priority: 100
timeout_in_minutes: 60
soft_fail: false
+36 -4
View File
@@ -5,7 +5,12 @@ set -ex
omni_source_dir=$(git rev-parse --show-toplevel)
base_image_name="xpu/vllm-omni-ci-base:${VLLM_VERSION:?VLLM_VERSION must be set}"
: "${VLLM_VERSION:?VLLM_VERSION must be set}"
# Official upstream XPU image. It is published from vllm's own
# docker/Dockerfile.xpu, which is exactly what build_vllm_base() clones and
# builds below, so pulling it is equivalent to building the fallback base.
upstream_base_image="vllm/vllm-openai-xpu:${VLLM_VERSION}"
local_base_image="xpu/vllm-omni-ci-base:${VLLM_VERSION}"
image_name="xpu/vllm-omni-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT must be set}"
container_name="xpu_${BUILDKITE_COMMIT}_$(
tr -dc A-Za-z0-9 </dev/urandom | head -c 10
@@ -49,9 +54,36 @@ build_vllm_base() (
"${vllm_source_dir}"
)
if [ -z "$(docker images -q "${base_image_name}")" ]; then
build_vllm_base
fi
# Resolve the vLLM base image, in priority order:
# 1. the published upstream image for VLLM_VERSION, pulled fresh (so moving
# tags like nightly are refreshed rather than served stale from disk);
# 2. a local copy of that same image, when the registry is unreachable;
# 3. VLLM_BASE, if that image is already on disk;
# 4. otherwise build VLLM_BASE from the matching vLLM source tag (~30min).
# VLLM_BASE names the fallback base and defaults to ${local_base_image}; set it
# to reuse or produce a base under a different tag.
resolve_vllm_base() {
if docker pull "${upstream_base_image}"; then
base_image_name="${upstream_base_image}"
return
fi
# Registry unreachable, but a previous run may have left the image behind.
if [ -n "$(docker images -q "${upstream_base_image}")" ]; then
echo "WARNING: could not pull ${upstream_base_image}; using local copy" >&2
base_image_name="${upstream_base_image}"
return
fi
echo "WARNING: ${upstream_base_image} unavailable remotely and locally" >&2
base_image_name="${VLLM_BASE:-${local_base_image}}"
if [ -z "$(docker images -q "${base_image_name}")" ]; then
echo "WARNING: building ${base_image_name} from vLLM ${VLLM_VERSION} source" >&2
build_vllm_base
fi
}
resolve_vllm_base
# Try building the docker image
docker_build "${image_name}" --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}"
+12
View File
@@ -10,6 +10,18 @@ steps:
-m 'core_model and npu and A3 and omni'
--run-level 'core_model'
- label: "Diffusion NPU layer tests"
key: ready-diffusion-npu-layer-tests
depends_on: upload-ready-pipeline
mirror_hardwares: a3_npu_2
commands:
- >-
pytest -s -v
tests/diffusion/layers/test_activation.py
tests/diffusion/layers/test_fused_qk_norm_rope_npu.py
-m 'core_model and npu and A3'
--run-level 'core_model'
- group: ":card_index_dividers: Omni Model Test"
key: ready-omni-group
depends_on: upload-ready-pipeline
+1 -1
View File
@@ -338,7 +338,7 @@ Classify the model's **CI priority** first:
**L2 & L3 online — same file, dual marks on the baseline smoke:** The **first / simplest** case in `test_{slug}.py` (default deploy, minimal steps, single prompt) should carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** function so L2 (`test-ready.yml`) and L3 (`test-merge.yml`) share one smoke test. Heavier deploy variants or API paths in the same file use **`advanced_model` only**. When L3 moves to nightly, migrate those heavier cases into `test_{slug}_expansion.py` with `full_model` and remove the dedicated `test-merge.yml` job (see `test_longcat_image_expansion.py`, `test_qwen_image_expansion.py`).
**L4 design (high priority):** Combine multiple supported features (Cache-DiT, TP, USP, CFG, HSDP, CPU offload, quantization) into **few parametrized** `OmniServerParams` rows so each feature appears in at least one case without exploding GPU jobs. Shard single-GPU vs multi-GPU cases across the nightly X2I/X2V function steps (1-GPU vs `distributed_cuda`).
**L4 design (high priority):** Combine multiple supported features (Cache-DiT, TP, USP, CFG, HSDP, CPU offload, quantization) into **few parametrized** `OmniServerParams` rows so each feature appears in at least one case without exploding GPU jobs. Shard single-GPU vs multi-GPU cases across the nightly X2I/X2V function steps (`cards_1` vs `not cards_1`).
**L4 design (medium / low):** One or two parametrized rows covering the best quality/perf trade-off; skip perf JSON unless the model is high priority.
@@ -34,7 +34,7 @@ but the changes are correct — stage the modified files and re-commit.
| `check-forbidden-imports` | Stdlib `re`/`base64`, pickle, Hugging Face Hub API, or direct Triton/TileLang | `import regex as re` and `pybase64`; use `vllm.transformers_utils.repo_utils`. Do **not** add the file to `allowed_files` without review |
| `check-torch-cuda-call` | New `torch.cuda.*` call site | Use `current_omni_platform` / `OmniPlatform`; do not grow `ALLOWED_FILES` without review |
| `check-tts-adapter-migration` | New `self._tts_model_type` branch in `serving_speech.py` | Put per-model logic in `tts_adapters/`. Lower `MAX_MODEL_TYPE_BRANCHES` when removing branches; do not raise it without review |
| `check-test-ci-coverage` | New `tests/**/test_*.py` missing level or hardware mark | Add `core_model`/`advanced_model`/… plus `cpu`/`cuda`/`hardware_test(` |
| `check-mark` | New `tests/**/test_*.py` missing level/hardware mark, or direct `pytest.mark.H100` | Add `core_model`/`advanced_model`/… plus `cpu`/`cuda`/`hardware_test(`; SKU marks only via helpers |
| `mypy-3.10` | Type error on changed `vllm_omni/` files | Fix the types; model trees are excluded. Extra versions: `pre-commit run --hook-stage manual mypy-3.12` |
| `markdownlint-cli2` | Docs/README markdown lint | Accept auto-fix under `docs/` / `recipes/` / root README |
| `check-buildkite` | Invalid `.buildkite/*.yml` | Fix the YAML; do not grow `SKIP_FILES` without review |
+3 -3
View File
@@ -52,8 +52,8 @@ Always attach markers deliberately:
- `tts` — speech synthesis / TTS-only models (`/v1/audio/speech`, voice clone, etc.)
- `diffusion` — generative diffusion models (image / audio / text / video from noise)
- **Cross-cutting area** (when relevant): `parallel`, `cache`, `example`, `benchmark`
- **Hardware**: `cpu`, `gpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `distributed_cuda`, …
- Optional: `slow`, distributed markers when multi-card is required
- **Hardware**: `cpu`, `gpu`, `cuda`, `rocm`, `npu`, …; SKU (`H100`, `L4`, …) and `cards_{n}` via `@hardware_test(...)` / `hardware_marks(...)` only — never `pytest.mark.H100` (`check-mark` rejects it)
- Optional: `slow`. Multi-card: `num_cards=N` on the helper (filter with `cards_N` / `not cards_1`)
**Baseline smoke (L2 + L3):** The simplest e2e case per model — default deploy, minimal request — should usually carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** test function so `test-ready.yml` and `test-merge.yml` share one test. `send_*_request` picks validation depth from `--run-level`. References: `test_voxcpm2_tts.py::test_text_to_audio_001`, `test_qwen3_tts_customvoice.py::test_text_to_audio_001`. Heavier scenarios use **`advanced_model` only**; L4 expansion uses **`full_model`**.
@@ -274,7 +274,7 @@ def test_text_to_video_001(omni_server, online_client) -> None:
**3. Cross-cutting rules**
- New `.py` files need the Omni SPDX header (`Copyright contributors to the vLLM-Omni project`, not `vLLM project`). `tests/` may keep stdlib `re`/`base64`; pickle is still banned unless the file is already on the pickle allowlist. Every collected `tests/**/test_*.py` needs a CI **level** mark and a **hardware** mark/helper (`check-test-ci-coverage` runs locally; GHA skips it). Do not grow `allowed_files` to land a test. Policy: [docs/contributing/README.md](../../../docs/contributing/README.md#linting).
- New `.py` files need the Omni SPDX header (`Copyright contributors to the vLLM-Omni project`, not `vLLM project`). `tests/` may keep stdlib `re`/`base64`; pickle is still banned unless the file is already on the pickle allowlist. Every collected `tests/**/test_*.py` needs a CI **level** mark and a **hardware** platform mark/helper (`check-mark` runs locally; GHA skips it). Do not write `pytest.mark.H100` / `pytest.mark.L4`; use the helpers so `cards_{n}` is attached. Do not grow `allowed_files` to land a test. Policy: [docs/contributing/README.md](../../../docs/contributing/README.md#linting).
- Reuse existing fixtures for the chosen scenario; do not mix “online client” assumptions into offline `OmniRunner` tests without a clear reason.
- Avoid external network dependency in assertions unless the scenario is explicitly “online serving” or doc examples that require a model hub (then align with CI secrets/cache).
- Keep **one test function = one intent** (one modality combo, one endpoint contract, or one acceleration combo).
@@ -79,7 +79,7 @@ E2E (L2+) should not rely on mocks unless documenting a rare exception; prefer r
- `full_model` — L4 nightly (`test-nightly.yml`); some expansion tests still carry both `advanced_model` and `full_model` during migration
2. **Model type** (pick one): `omni`, `tts`, or `diffusion`
3. **Cross-cutting** when relevant: `parallel`, `cache`, `example`, `benchmark`
4. **Hardware**: `cpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `distributed_cuda`, …
4. **Hardware**: `cpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `cards_1`, …
5. Multi-card: `@hardware_test(...)` in `tests/helpers/mark.py`
## Command Templates
@@ -244,7 +244,7 @@ pytest -s -v dfx/reliability/invalid_param_test/test_invalid_image_generation.py
```bash
cd tests
pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
pytest -s -v -m "core_model and L4 and not cards_1" --run-level=core_model
```
### Concrete e2e paths (common in-tree)
+1 -1
View File
@@ -27,4 +27,4 @@ jobs:
# Don't run historical-debt hooks against the whole tree in CI yet.
# They still run locally on changed files.
env:
SKIP: check-test-ci-coverage,markdownlint-cli2,shellcheck,check-spdx-header,mypy-3.10
SKIP: check-mark,markdownlint-cli2,shellcheck,check-spdx-header,mypy-3.10
+2 -2
View File
@@ -105,8 +105,8 @@ repos:
<<: *mypy_common
stages: [manual]
- id: check-test-ci-coverage
name: Ensure test files have pytestmarks
- id: check-mark
name: Ensure test files have CI marks (no direct SKU)
entry: python tools/pre_commit/check_test_marks.py
language: python
types: [python]
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.27.0
ARG VLLM_ASCEND_TAG=v0.28.0
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
# WORKDIR /vllm-workspace/vllm
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.27.0-a3
ARG VLLM_ASCEND_TAG=v0.28.0-a3
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
# WORKDIR /vllm-workspace/vllm
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.27.0
ARG VLLM_ASCEND_TAG=v0.28.0
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
ARG APP_DIR=/vllm-workspace/vllm-omni
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.27.0-a3
ARG VLLM_ASCEND_TAG=v0.28.0-a3
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
ARG APP_DIR=/vllm-workspace/vllm-omni
+27 -134
View File
@@ -1,140 +1,33 @@
# Argument to configure vllm base image if pre-built
ARG VLLM_BASE=vllm-base
FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base
WORKDIR /workspace/
ARG PYTHON_VERSION=3.12
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
RUN apt clean && apt-get update -y && \
apt-get install -y --no-install-recommends --fix-missing \
curl \
ffmpeg \
git \
libsndfile1 \
libsm6 \
libxext6 \
libgl1 \
lsb-release \
libaio-dev \
numactl \
wget \
vim \
python3.12 \
python3.12-dev \
python3-pip
# Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer.
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
printf '%s\n' \
'Package: intel-oneapi-* intel-deep-learning-essentials* intel-pti*' \
'Pin: version 2025.3*' \
'Pin-Priority: 1001' \
> /etc/apt/preferences.d/oneapi-2025.3.pref && \
apt-get update -y && \
apt-get install -y --no-install-recommends \
intel-oneapi-compiler-dpcpp-cpp-2025.3 \
intel-oneapi-mkl-devel-2025.3 \
intel-oneapi-dnnl-devel-2025.3 && \
rm -rf /var/lib/apt/lists/*
# Install UMD with fixed version
RUN mkdir neo && \
cd neo && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
dpkg -i *.deb && \
cd .. && \
rm -rf neo
ENV PATH="/root/.local/bin:$PATH"
ENV VIRTUAL_ENV="/opt/venv"
ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
RUN curl -LsSf https://astral.sh/uv/install.sh | sh
RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh"
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \
bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
rm "${ONECCL_INSTALLER}" && \
echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc
RUN rm -f /opt/intel/oneapi/ccl/latest && \
ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest
SHELL ["bash", "-c"]
CMD ["bash", "-c", "source /root/.bashrc && exec bash"]
WORKDIR /workspace/
ENV UV_HTTP_TIMEOUT=500
# Configure package index for XPU
ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
ENV UV_INDEX_STRATEGY="unsafe-best-match"
ENV UV_LINK_MODE="copy"
ARG VLLM_VERSION=v0.27.0
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
WORKDIR /workspace/vllm
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --upgrade pip && \
uv pip install -r requirements/xpu.txt && \
uv pip install setuptools_scm &&\
uv pip install grpcio-tools protobuf nanobind && \
source /opt/intel/oneapi/setvars.sh --force && \
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}"
ENV LD_LIBRARY_PATH="/opt/intel/oneapi/ccl/2021.15/lib:$LD_LIBRARY_PATH:/usr/local/lib/"
ENV VLLM_TARGET_DEVICE=xpu
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip uninstall triton triton-xpu && \
uv pip install triton-xpu==3.7.2 && \
uv pip install --no-build-isolation --no-deps .
CMD ["/bin/bash"]
FROM vllm-base AS vllm-openai
# install nixl from source code
RUN uv pip install nixl
# install additional dependencies for openai api server
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install accelerate hf_transfer pytest pytest_asyncio lm_eval[api] modelscope
# install development dependencies (for testing)
RUN uv pip install -e tests/vllm_test_utils
# ensure vllm is properly installed
RUN python -c "import vllm, inspect; print(vllm.__file__)"
RUN uv pip show vllm
CMD ["/bin/bash"]
ENTRYPOINT []
# Target vLLM version: v0.28.0. Keep VLLM_VERSION aligned with the target vLLM
# release on every rebase.
#
# This file layers vLLM-Omni on top of the published upstream XPU image, which
# already ships a self-consistent XPU runtime (torch-xpu, oneCCL/UCX+NIXL, the
# pinned UMD/compute-runtime stack, and vLLM itself) built from
# https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu. Docker
# pulls that base automatically, so there is nothing to build here beyond the
# Omni layers. Mirrors how docker/Dockerfile.ci consumes vllm/vllm-openai.
#
# If the tag does not exist yet (targeting an unreleased vLLM commit), build the
# base from upstream's own Dockerfile and point VLLM_BASE at the result:
#
# git clone --depth 1 --branch $VLLM_VERSION https://github.com/vllm-project/vllm /tmp/vllm
# docker build -t vllm-openai-xpu:local --target vllm-openai \
# -f /tmp/vllm/docker/Dockerfile.xpu /tmp/vllm
# docker build -f docker/Dockerfile.xpu --build-arg VLLM_BASE=vllm-openai-xpu:local .
#
# .buildkite/intel/scripts/run-xpu-test.sh automates that pull-then-build fallback.
ARG VLLM_BASE_IMAGE=vllm/vllm-openai-xpu
ARG VLLM_VERSION=v0.28.0
ARG VLLM_BASE=${VLLM_BASE_IMAGE}:${VLLM_VERSION}
FROM ${VLLM_BASE} AS vllm-omni
# The base image ships ENTRYPOINT ["vllm", "serve"]; clear it so this stage stays
# usable as a plain shell/test container regardless of which VLLM_BASE was used.
ENTRYPOINT []
CMD ["/bin/bash"]
RUN apt-get update && \
apt-get install -y espeak-ng ffmpeg git jq && \
apt-get clean && \
+8 -6
View File
@@ -49,7 +49,7 @@ vLLM-Omni's `pre-commit` hooks will now run automatically every time you commit.
`--all-files` run stays green while historical debt is cleaned up. The
current `SKIP` list in
[`.github/workflows/pre-commit.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.github/workflows/pre-commit.yml)
is `check-test-ci-coverage`, `markdownlint-cli2`, `shellcheck`,
is `check-mark`, `markdownlint-cli2`, `shellcheck`,
`check-spdx-header`, and `mypy-3.10`. Those hooks still run on **your
commit** for changed files. A passing GitHub pre-commit check does not mean
they passed locally. Hooks **not** on that list (forbidden imports,
@@ -66,7 +66,7 @@ already there). `check-pickle-imports` is gone: pickle is now one rule inside
| `markdownlint-cli2` | Markdown in `docs/`, `recipes/`, `README.md`, `CONTRIBUTING.md` (not `.claude/` / `.cursor/` / `CLAUDE.md`) | skipped |
| `mypy-3.10` | Type-check changed `vllm_omni/` files (model trees excluded). `tests/` uses `--follow-imports skip` | skipped |
| `mypy-3.11` / `3.12` / `3.13` | Same checker, extra Python versions | not installed; `pre-commit run --hook-stage manual mypy-3.12` |
| `check-test-ci-coverage` | Every `tests/**/test_*.py` has a CI level mark and a hardware mark/helper | skipped |
| `check-mark` | Every `tests/**/test_*.py` has a CI level mark and a hardware platform mark or helper; no direct `pytest.mark.H100` / SKU | skipped |
| `check-tts-adapter-migration` | `self._tts_model_type` branches in `serving_speech.py` must not increase | runs |
| `shellcheck` | `*.sh` quoting / undefined vars | skipped |
| `check-spdx-header` | Omni SPDX header on `.py` / `.pyi` / `.sh` / `.rs` / `.proto` | skipped |
@@ -144,11 +144,13 @@ pre-commit run --hook-stage manual mypy-3.12
#### Test CI marks
`check-test-ci-coverage` requires each collected test module under `tests/` to
`check-mark` requires each collected test module under `tests/` to
have at least one CI **level** mark (`core_model`, `advanced_model`,
`full_model`, `local_model`, or `slow`) and a **hardware** mark (`cpu`, `cuda`,
`H100`, …) or helper (`hardware_test(` / `hardware_marks(`). GitHub Actions
skips this hook; local commit does not. See the
`full_model`, `local_model`, or `slow`) and a **hardware** platform mark
(`cpu`, `cuda`, …) or helper (`hardware_test(` / `hardware_marks(`). SKU marks
(`H100`, `L4`, …) must come from those helpers so `cards_{n}` is attached;
direct `pytest.mark.H100` is rejected. GitHub Actions skips this hook; local
commit does not. See the
[test writing guide](./ci/test_writing_guide.md).
#### TTS adapter ratchet
+1 -1
View File
@@ -402,7 +402,7 @@ loads the model once per mode.
Before splitting, confirm each half actually collects at least one test under the job's forwarded filters — pytest exits with code 5 ("no tests collected") if a half is empty, which the script reports as a failure, red-ing a job that used to pass as one combined invocation.
Also check the tests' `num_cards` against the job's `mirror_hardwares` preset.
`cuda_marks` turns `num_cards > 1` into `skipif(device_count() < num_cards)`, so a
`_cuda_marks` turns `num_cards > 1` into `skipif(device_count() < num_cards)`, so a
test asking for more GPUs than the preset provides is silently skipped and its
coverage artifact comes back empty. Splitting such a job produces a file that
looks fine and measures nothing — "Diffusion · Z Image Test" is the current
@@ -81,8 +81,8 @@ Supported form:
| Form | Example | Effect |
| ---- | ------- | ------ |
| Array (single platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "diffusion"]` | Hardware marks from `hardware_marks(...)`; `num_cards` > 1 adds `distributed_*` (+ CUDA `skipif_cuda` when applicable). String entries become extra pytest marks. |
| Array (multi-platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100", "rocm": "MI325", "npu": "A2"}, "num_cards": {"cuda": 2, "rocm": 2, "npu": 1}}}, "full_model", "omni"]` | Same as above, but declares **multiple platforms** in one case. Filter examples: `-m "full_model and H100 and cuda"`, `-m "full_model and MI325 and rocm"`. |
| Array (single platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "diffusion"]` | Hardware marks from `hardware_marks(...)`; `num_cards` adds `cards_{n}` (including `cards_1` when omitted). String entries become extra pytest marks. |
| Array (multi-platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100", "rocm": "MI325"}, "num_cards": 2}}, "full_model", "omni"]` | Same as above, but declares **multiple platforms** on one item. All platforms must share the same `num_cards` (int or identical dict values). Different counts need `@hardware_test` or one `hardware_marks()` per `pytest.param`. Filter examples: `-m "full_model and H100 and cards_2 and cuda"`, `-m "full_model and MI325 and cards_2 and rocm"`. |
Recommended for L4 perf cases:
@@ -157,7 +157,7 @@ pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
```
<!-- markdownlint-disable-next-line MD051 -->
See also [Markers for Tests](#markers-for-tests) for registered hardware markers (`H100`, `L4`, `cuda`, `distributed_cuda`, …).
See also [Markers for Tests](#markers-for-tests) for registered hardware markers (`H100`, `L4`, `cuda`, `cards_1`, …).
#### `server_params` Configuration
+4 -4
View File
@@ -88,7 +88,7 @@ Failed jobs: 1/2
```bash
cd tests
pytest -s -v test_xxxx.py --run-level=core_model
pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
pytest -s -v -m "core_model and L4 and not cards_1" --run-level=core_model
```
=== "L3 level"
@@ -119,7 +119,7 @@ Failed jobs: 1/2
```bash
pytest -s -v test_xxxx.py --run-level=advanced_model
pytest -s -v -m "advanced_model and distributed_cuda and L4" --run-level=advanced_model
pytest -s -v -m "advanced_model and L4 and not cards_1" --run-level=advanced_model
```
=== "L4 level"
@@ -153,12 +153,12 @@ Failed jobs: 1/2
```bash
cd tests
pytest -s -v test_xxxx.py --run-level=full_model
pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model
pytest -s -v -m "full_model and L4 and not cards_1" --run-level=full_model
pytest -s -v -m "full_model and (omni or tts) and H100" --run-level=full_model
```
If you only want to run specific test cases on a particular platform, you can use:
```bash
pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model
pytest -s -v -m "full_model and L4 and not cards_1" --run-level=full_model
```
Note: ``run_benchmark.py`` and ``run_diffusion_benchmark.py`` accept an optional ``--test-config-file``. If omitted, each loads every ``*.json`` under ``tests/dfx/perf/tests/`` (omni/tts vs diffusion split by ``is_diffusion_perf_config``) and pytest ``-m`` filters by each case's JSON ``mark``:
```bash
+34 -47
View File
@@ -48,17 +48,12 @@ Defined in `pyproject.toml`:
| `MI325` | Tests that require MI325 GPU (AMD/ROCm) * |
| `A2` | Tests that require A2 NPU * |
| `A3` | Tests that require A3 NPU * |
| `distributed_cuda` | Tests that require multi cards on CUDA platform * |
| `distributed_rocm` | Tests that require multi cards on ROCm platform * |
| `distributed_npu` | Tests that require multi cards on NPU platform * |
| `skipif_cuda` | Skip if the num of CUDA cards is less than the required * |
| `skipif_rocm` | Skip if the num of ROCm cards is less than the required * |
| `skipif_npu` | Skip if the num of NPU cards is less than the required * |
| `cards_{n}` | Tests that require *n* accelerator cards (`cards_1` … `cards_8`); auto-added from `num_cards` * |
| `slow` | Slow tests (may skip in quick CI) |
| `benchmark` | Benchmark tests (decorator on runner test functions; perf JSON uses `full_model` + type marker instead) |
| `local_model` | Tests requiring local / non-HF-hub model weights |
\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility).
\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility). Do **not** write `pytest.mark.H100` / `pytest.mark.L4` (or other SKU names) by hand: `check-mark` rejects that so `cards_{n}` cannot be skipped.
#### Example usage for markers
@@ -81,25 +76,23 @@ def test_video_to_audio()
This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/helpers/mark.py` performs the following actions:
1. **Applies platform and resource markers**
Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `B60`, `A2`, `A3`).
Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `H200`, `B200`, `MI325`, `B60`, `A2`, `A3`).
```python
@pytest.mark.cuda
@pytest.mark.L4
```
2. **Handles multi-card (distributed) scenarios**
For tests requiring multiple cards, it automatically adds distributed markers such as `distributed_cuda`, `distributed_rocm`, or `distributed_npu`.
```python
@pytest.mark.distributed_cuda(num_cards=num_cards)
```
2. **Handles card-count filtering**
Adds a filterable marker `cards_{n}` (including `cards_1` when `num_cards` is omitted / 1). Use this with the SKU and platform, e.g. `-m "H100 and cards_4 and cuda"` vs `-m "H100 and cards_2 and cuda"`. Any multi-card CUDA test is `-m "cuda and not cards_1"`.
When `res` lists more than one platform **and** `num_cards` differs by platform, `@hardware_test` collects **one variant per platform** so those `cards_{n}` marks do not stack. `-m "H100 and cards_4 and cuda"` then selects only the CUDA/H100 4-card variant, not a ROCm 2-card sibling. The same count on every platform stays a single item.
3. **Supports flexible card requirements**
Accepts `num_cards` as either a single integer for all platforms or as a dictionary with per-platform values. If not specified, defaults to 1 card per platform.
4. **Integrates resource validation**
On CUDA, adds a skip marker (`skipif_cuda`) if the system does not have the required number of devices.
Support for `skipif_rocm` and `skipif_npu` will be implemented later.
On CUDA and XPU, when `num_cards > 1`, attaches pytest's built-in `skipif` if the machine has fewer devices than requested. This is not a named `skipif_*` marker and is not used for `-m` filtering.
5. **Works with pytest filtering**
Allows tests to be filtered and selected at runtime using standard pytest marker expressions (e.g., `-m "distributed_cuda and L4"`).
Allows tests to be filtered and selected at runtime using standard pytest marker expressions (e.g., `-m "H100 and cards_2"`, `-m "L4 and not cards_1"`).
##### Example usage for decorator
@@ -117,13 +110,14 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
num_cards=2,
)
```
- `res` must be a dict; supported resources: CUDA (L4/H100), ROCm (MI325), XPU (B60), MUSA (S5000), NPU (A2/A3)
- `res` must be a dict. SKU values come from `pyproject.toml` markers tagged `[hardware-resource]` plus the platform key (`[cuda]`, `[rocm]`, `[xpu]`, `[npu]`, `[musa]`).
- `num_cards` can be int (all platforms) or dict (per platform); defaults to 1 when missing
- Distributed markers (`distributed_cuda`, `distributed_rocm`, `distributed_npu`) are auto-added for multi-card cases
- Filtering examples:
- CUDA only: `pytest -m "distributed_cuda and L4"`
- ROCm only: `pytest -m "distributed_rocm and MI325"`
- NPU only: `pytest -m "distributed_npu"`
- CUDA H100 4-card: `pytest -m "H100 and cards_4 and cuda"`
- CUDA H100 2-card: `pytest -m "H100 and cards_2 and cuda"`
- CUDA multi-card (any count): `pytest -m "L4 and not cards_1"`
- ROCm multi-card: `pytest -m "MI325 and not cards_1"`
- NPU multi-card: `pytest -m "npu and not cards_1"`
#### Function: `hardware_marks`
@@ -179,37 +173,30 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
[tool.pytest.ini_options]
markers = [
# ... existing markers ...
"tpu: Tests that require TPU device",
"TPU_V3: Tests that require TPU v3 hardware",
"distributed_tpu: Tests that require multiple TPU devices",
"tpu: [hardware-platform] Tests that require TPU device",
# Add [gpu] next to [hardware-platform] if this platform should also get pytest.mark.gpu.
"TPU_V3: [hardware-resource] [tpu] Tests that require TPU v3 hardware",
"cards_1: [hardware-cards] Tests that require 1 accelerator card",
"cards_2: [hardware-cards] Tests that require 2 accelerator cards",
]
```
`check-mark` treats `[hardware-platform]` names (including `cpu`/`gpu`) as allowed `pytest.mark.<name>` values. `hardware_test(res=...)` still rejects `cpu`/`gpu` as `res` keys.
2. **Implement a marker construction function for your platform** in `vllm-omni/tests/helpers/mark.py`:
```python
# In vllm-omni/tests/helpers/mark.py
def tpu_marks(*, res: str, num_cards: int):
test_platform = pytest.mark.tpu
if res == "TPU_V3":
test_resource = pytest.mark.TPU_V3
else:
raise ValueError(
f"Invalid TPU resource type: {res}. Supported: TPU_V3")
if num_cards == 1:
return [test_platform, test_resource]
else:
test_distributed = pytest.mark.distributed_tpu(num_cards=num_cards)
# Optionally: add skipif_tpu when implemented
return [test_platform, test_resource, test_distributed]
def _tpu_marks(*, res: str, num_cards: int):
_require_sku("tpu", res)
return [pytest.mark.tpu, getattr(pytest.mark, res), _cards_mark(num_cards)]
```
3. **Update `hardware_marks` to recognize your new platform**:
In the relevant place (see the `hardware_marks` implementation), add:
3. **Register the builder** in `_marks_for_platform`:
```python
if platform == "tpu":
marks = tpu_marks(res=resource, num_cards=cards)
builders = {
# ... existing platforms ...
"tpu": _tpu_marks,
}
```
(`hardware_test` calls `hardware_marks` internally, so both will pick up the change.)
(`hardware_test` / `hardware_marks` both go through `_marks_for_platform`.)
4. **(Recommended) Add a test using your new markers**:
```python
@hardware_test(
@@ -223,11 +210,11 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
**Summary**:
- Add pytest markers for your new platform/resources
- Implement a marker function (`xxx_marks`)
- Plug into `hardware_marks`
- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers, distribution, and isolation!
- Implement a private marker function (`_xxx_marks`)
- Register it in `_marks_for_platform`
- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers and `cards_{n}` filtering.
See code in `vllm-omni/tests/helpers/mark.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).
See code in `vllm-omni/tests/helpers/mark.py` for existing examples (`_cuda_marks`, `_rocm_marks`, `_npu_marks`).
## Test case style
+2 -2
View File
@@ -438,7 +438,7 @@ selects the scheduler and worker family:
Key topology fields belong to `StagePipelineConfig`:
| Field | Description |
|-------|-------------|
| ------- | ------------- |
| `model_stage` | Logical stage name (`ar_stage`, `decoder`, etc.). |
| `execution_type` | Runtime family for this stage. |
| `input_sources` | Upstream stage IDs that provide input. |
@@ -960,7 +960,7 @@ Common failures and fixes:
| `check-forbidden-imports` | Stdlib `re`/`base64`, pickle, Hugging Face Hub API, or direct Triton/TileLang | `import regex as re` and `pybase64`; do not grow the allowlist without review |
| `check-torch-cuda-call` | New `torch.cuda.*` call site | Use `current_omni_platform`; do not grow `ALLOWED_FILES` without review |
| `check-tts-adapter-migration` | New `_tts_model_type` branch in `serving_speech.py` | Put logic in `tts_adapters/`; lower `MAX_MODEL_TYPE_BRANCHES` when removing branches |
| `check-test-ci-coverage` | Test file missing level or hardware mark | Add `core_model`/… plus `cpu`/`cuda`/`hardware_test(` |
| `check-mark` | Test file missing level/hardware mark, or direct `pytest.mark.H100` | Add `core_model`/… plus `cpu`/`cuda`/`hardware_test(`; SKU marks only via helpers |
| `shellcheck` | Native `shellcheck` missing, or script warning | Install via apt/dnf/brew (or `shellcheck.exe` on Windows). See [Linting](../README.md#linting) |
| `mypy-3.10` / `markdownlint-cli2` / `check-buildkite` | Types, docs markdown, or Buildkite YAML | See [Linting](../README.md#linting) |
@@ -26,6 +26,8 @@ Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsu
| **TP > 1** | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. |
| **HSDP** | ❌ Rejected to avoid double-sharding parameters. | ⚠️ Limited end-to-end coverage. |
| **Per-tensor online FP8 linears** | ✅ Ordinary loader finalizes weights and scales before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
| **Online INT8 linears** | ✅ Ordinary loader finalizes the int8 weight (contiguous or transposed view) and fp32 scale before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
| **Online MXFP8 linears** | ✅ Ordinary loader finalizes the fp8 weight (contiguous) and e8m0 block scale (contiguous or transposed view) before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
| **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. |
| **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. |
| **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. |
@@ -438,11 +440,15 @@ requirements.
Direct checkpoint mmap can back either transfer path. It is currently limited
to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the
ordinary loader finalizes their runtime weights and scales; DLO then shards and
reconstructs those tensors with their recorded layouts. Other online methods
must use `--dlo-no-use-allgather` or disable online quantization until their
runtime layouts are validated.
ordinary loader. Per-tensor online FP8, online INT8, and online MXFP8 linears
can use DLO AllGather after the ordinary loader finalizes their runtime
weights and scales; DLO then shards and reconstructs those tensors with their
recorded layouts (they all keep plain 1-byte dtypes over ordinary strided
views: online INT8 has an int8 weight plus fp32 scale, online MXFP8 has an
fp8 weight plus e8m0 block scale, each with a contiguous or transposed-view
finalized tensor, all covered by the same physical-order packing as online
FP8). Other online methods must use `--dlo-no-use-allgather` or disable
online quantization until their runtime layouts are validated.
The Host Weight Runtime representation, publication, and no-AllGather consumer
contracts are merged; see
@@ -459,8 +465,8 @@ Current source-level validation includes:
- HSDP + DLO without AllGather acceptance at configuration level;
- loader preflight fallback for TP, HSDP, online quantization, unknown custom
loaders, missing keys, and shape/dtype mismatches;
- ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
sharding of finalized weights and scales;
- ordinary-loader fallback for per-tensor online FP8, online INT8, and online
MXFP8 linears followed by DLO sharding of finalized weights and scales;
- exact loader-to-backend plan transfer and ordinary-loader fallback;
- ordered publication hashing, overlapped durability, unordered parallel
checksum fallback, and node-local source-digest reuse/invalidation;
+69 -49
View File
@@ -22,59 +22,60 @@ We use **Qwen-Image** as the reference implementation for decode parallel, and *
**VAE Patch Parallelism** is an acceleration technique for both **encoding** and **decoding**. Instead of processing the entire tensor at once, the tensor is:
+ Split into multiple spatial tiles
- Split into multiple spatial tiles
+ Distributed across multiple ranks
- Distributed across multiple ranks
+ Encoded/Decoded in parallel
- Encoded/Decoded in parallel
+ Merged to reconstruct the final output
- Merged to reconstruct the final output
This approach:
+ Distributes computation across multiple devices
- Distributes computation across multiple devices
+ Reduces peak memory usage per device
- Reduces peak memory usage per device
+ Accelerates encoding/decoding latency
- Accelerates encoding/decoding latency
### When to Use Encode vs Decode Parallel
| Operation | Use Case | Example |
|-----------|----------|---------|
| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |
| Operation | Use Case | Example |
| ------------------- | ---------------------------- | --------------------------------- |
| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |
### Architecture
We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/decoding.
The executor is model-agnostic and accepts three function parameters:
+ split – Partition the latent into tiles
- split – Partition the latent into tiles
+ exec – Decode a single tile
- exec – Decode a single tile
+ merge – Combine decoded tiles into the final output
- merge – Combine decoded tiles into the final output
#### Execution Flow
+ Call split(z) to generate a list of TileTask and a GridSpec
- Call split(z) to generate a list of TileTask and a GridSpec
+ Dispatch tasks across ranks using workload-based balancing
- Dispatch tasks across ranks using workload-based balancing
+ Each rank executes exec(task) on its assigned tiles
- Each rank executes exec(task) on its assigned tiles
+ Gather decoded tile results to rank 0
- Gather decoded tile results to rank 0
+ Rank 0 performs merge(...)
- Rank 0 performs merge(...)
+ (Optional) Broadcast final result to all ranks
- (Optional) Broadcast final result to all ranks
This design separates:
+ Distributed execution logic
- Distributed execution logic
+ Model-specific tiling and merging logic
- Model-specific tiling and merging logic
#### Why split / exec / merge is necessary?
@@ -82,23 +83,23 @@ The latent tensor cannot be arbitrarily partitioned.
During decoding:
+ Each output pixel may depend on neighboring pixels
- Each output pixel may depend on neighboring pixels
+ The receptive field is model-dependent
- The receptive field is model-dependent
Therefore:
+ Tiles must include overlap
- Tiles must include overlap
+ Merge must perform blending to avoid seams
- Merge must perform blending to avoid seams
## Step-by-Step Implementation (Decode)
### Step 1: Implement DistributedAutoencoderKLQwenImage
`QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
```
```python
class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin):
@classmethod
def from_pretrained(cls, *args: Any, **kwargs: Any):
@@ -106,28 +107,35 @@ class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMi
model.init_distributed()
return model
```
**Key points**:
+ Inherit both AutoencoderKLQwenImage and DistributedVaeMixin
+ Call init_distributed() after loading weights
- Inherit both AutoencoderKLQwenImage and DistributedVaeMixin
- Call init_distributed() after loading weights
### Step 2: Implement split/exec/merge
Reuse `AutoencoderKLQwenImage.tiled_decode` logic and divide it into three stages. And we need return tiles with `GridSpec` and `TileTask`:
```
```python
class GridSpec:
split_dims: tuple[int, ...] # Tensor dimensions being split (e.g., (2, 3) for (B, C, H, W))
grid_shape: tuple[int, ...] # Tile grid layout (num_rows, num_cols)
tile_spec: dict = field(default_factory=dict) # Metadata required for merging
output_dtype: torch.dtype | None = None # Final output dtype
```
```
```python
class TileTask:
tile_id: int # task id
grid_coord: tuple[int, ...] # Tile position in grid
tensor: torch.Tensor | list[torch.Tensor] # The tile tensor
workload: int | float = 1 # Used for load balancing (e.g., tile area)
```
And tiled base split/exec/merge as follow:
```
```python
def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
# mostly copy from AutoencoderKL
_, _, num_frames, height, width = z.shape
@@ -210,11 +218,14 @@ def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid
```
### Step 3: Override tiled_decode
We need to override tiled_decode, the main logic is:
+ check distributed is enabled
+ select split/exec/merge
+ Invoke self.distributed_executor.execute to decode
```
- check distributed is enabled
- select split/exec/merge
- Invoke self.distributed_executor.execute to decode
```python
def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
if not self.is_distributed_enabled():
return super().tiled_decode(z, return_dict=return_dict)
@@ -230,11 +241,14 @@ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
return DecoderOutput(sample=result)
```
`broadcast_result` is set to True or False depending on the model; when enabled, the result will be used even on ranks other than 0.
### Step 4: Modify Pipeline
Change vae model from AutoencoderKLQwenImage to DistributedAutoencoderKLQwenImage
```
```diff
class YourModelPipeline(nn.Module):
def __init__(
self,
@@ -259,8 +273,8 @@ For models that require VAE encoding (e.g., Image-to-Video), you can also parall
Similar to decode, split the input tensor into tiles. Key considerations:
+ **Patchify handling**: If the model uses `patch_size`, scale tile parameters accordingly
+ **Temporal chunking**: Video VAEs may have temporal compression (e.g., 4x)
- **Patchify handling**: If the model uses `patch_size`, scale tile parameters accordingly
- **Temporal chunking**: Video VAEs may have temporal compression (e.g., 4x)
```python
def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
@@ -406,23 +420,28 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
**Key differences from decode parallel:**
| Aspect | Decode Parallel | Encode Parallel |
|--------|-----------------|-----------------|
| -------- | ----------------- | ----------------- |
| `broadcast_result` | Often `False` (only rank 0 needs output) | `True` (all ranks need latents for diffusion) |
| Patchify | Applied in merge (unpatchify) | Handled by parent `_encode()` before `tiled_encode()` |
| Temporal chunking | Frame-by-frame | Chunk-based (e.g., 1 + 4n frames) |
## Testing
Verify numerical consistency between:
+ vae_patch_parallel_size = 1
+ vae_patch_parallel_size = N
Verify numerical consistency between:
- vae_patch_parallel_size = 1
- vae_patch_parallel_size = N
Example:
```python
torch.allclose(output_1, output_n, atol=1e-5)
```
Testing requirements:
+ Fix random seed
+ Use identical tiling strategy
- Fix random seed
- Use identical tiling strategy
```python
m = Omni(
@@ -434,6 +453,7 @@ m = Omni(
),
)
```
When vae_patch_parallel_size is larger than the DiT world size, it will automatically fall back to using the DiT world size instead.
## Reference Implementations
@@ -441,7 +461,7 @@ When vae_patch_parallel_size is larger than the DiT world size, it will automati
Complete examples in the codebase:
| Model | Path | Decode Parallel | Encode Parallel |
|-------|------|-----------------|-----------------|
| ------- | ------ | ----------------- | ----------------- |
| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | ✅ | ❌ |
| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | ✅ | ✅ |
| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | ✅ | ❌ |
@@ -456,7 +476,7 @@ The tile-parallel executor above assigns independent spatial **tiles** to ranks.
### How it differs from tile parallel
| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
|--------|--------------------------|-----------------------------------------------|
| -------- | -------------------------- | ----------------------------------------------- |
| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W |
| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) |
| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
@@ -481,7 +501,7 @@ serve.py (--vae-parallel-mode) / OmniEngineArgs
### Notes
- The decoder is patched **in place** the first time spatial-shard decode runs and is bound to a single split dimension for the lifetime of the VAE instance.
- Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `distributed_cuda`).
- Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `cards_2`).
---
@@ -37,6 +37,7 @@ validation_paths:
- tests/engine/test_async_omni_engine_input.py
- tests/engine/test_async_omni_engine_outputs.py
- tests/engine/test_async_omni_engine_abort.py
- tests/e2e/offline_inference/test_qwen3_omni.py
- tests/engine/test_async_omni_engine_stage_init.py
- tests/engine/test_orchestrator.py
- tests/engine/test_orchestrator_error_handling.py
@@ -113,6 +114,15 @@ Numbers become append-only after normative promotion.
Test routing, output ordering, queue correlation, cancellation, failure
propagation, shutdown ordering, and representative multi-stage execution.
AR abort of a final-stage LLM request must deliver a terminal output with
`finish_reason="abort"` and the generated prefix so collocated training can
resume; diffusion abort remains whole-sample retry. Aborting a parallel-sampling
child must drop `parent_requests` once no children remain. EngineCore control
RPCs (`sleep` / `wake_up` / `pause_scheduler` / `resume_scheduler`) must
propagate worker exceptions rather than returning `{"supported": False}`.
When a request has multiple final output stages, only the last abort message
is request-terminal. Output-processor abort state is committed only after the
physical EngineCore abort succeeds.
Changes that cross into stage runtime or public error behavior require review
from that contract's owners.
+8 -1
View File
@@ -43,6 +43,7 @@ depends_on:
validation_paths:
- tests/entrypoints/test_omni_entrypoints.py
- tests/entrypoints/test_async_omni.py
- tests/entrypoints/test_async_omni_pause_sleep_routing.py
- tests/entrypoints/test_async_omni_duplex.py
- tests/entrypoints/test_serve.py
- tests/entrypoints/test_stream_finish_reason.py
@@ -119,7 +120,13 @@ promotion.
Test request validation, protocol conversion, model-adapter routing,
streaming, session identity, disconnect/cancellation, and error mapping for
each affected offline or serving entrypoint.
each affected offline or serving entrypoint. Sleep must wait for in-flight
`generate()` admission before EngineCore offload; `wake_up` does not resume
admission — callers must `resume_generation()`. Sleeping tags are tracked per
stage so `wake_up(stage_ids=[0])` does not skip a later `wake_up(stage_ids=[1])`.
Streaming input pumps take an admission slot immediately before each EngineCore
ADD or update, not while waiting for the next client chunk. Frontend abort
keeps `request_states` until `generate()` consumes the terminal output.
## Promotion gate
+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
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.
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.
@@ -23,6 +23,19 @@ vLLM-Omni currently recommends using the Docker image setup steps below.
DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g .
```
This layers vLLM-Omni on top of the published `vllm/vllm-openai-xpu:<VLLM_VERSION>`
base image, which Docker pulls automatically. To target a different vLLM release,
pass `--build-arg VLLM_VERSION=<tag>`. If that tag has not been published yet, build
the base from upstream's own Dockerfile first and point the build at it:
```bash
git clone --depth 1 --branch "$VLLM_VERSION" https://github.com/vllm-project/vllm /tmp/vllm
DOCKER_BUILDKIT=1 docker build -t vllm-openai-xpu:local --target vllm-openai \
-f /tmp/vllm/docker/Dockerfile.xpu /tmp/vllm
DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g \
--build-arg VLLM_BASE=vllm-openai-xpu:local .
```
#### Launch the docker image
##### Launch with OpenAI API Server
+11 -11
View File
@@ -26,15 +26,15 @@ th {
| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | | — |
| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | | — |
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) |
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) |
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) |
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) |
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | — |
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | — |
| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | ✅︎ | — |
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — |
@@ -42,7 +42,7 @@ th {
| `SanaWmPipeline` | SANA-WM | `BBBBruce/SANA-WM_bidirectional-stage1-diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-WM.md) |
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | — |
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | ✅︎ | — |
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
| `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — |
| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | | — |
@@ -62,12 +62,12 @@ th {
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `LongCatVideoAvatarPipeline` | LongCat-Video-Avatar-1.5 A2V/AI2V (native single-speaker and multi-speaker AI2V/AVC) | `meituan-longcat/LongCat-Video-Avatar-1.5` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md) |
| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | | — |
| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | ✅︎ | — |
| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — |
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — |
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — |
| `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` | ✅︎ | | | | — |
@@ -78,10 +78,10 @@ th {
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — |
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — |
| `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | | — |
| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | ✅︎ | — |
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | | — |
| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | | — |
| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | | — |
@@ -101,10 +101,10 @@ th {
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — |
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — |
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | — |
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | ✅︎ | — |
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | | — |
| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | ✅︎ | — |
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
| `HiDreamO1ImagePipeline` | HiDream-O1-Image | `HiDream-ai/HiDream-O1-Image` | ✅︎ | | | | — |
@@ -17,6 +17,8 @@ globally.
| --- | --- | --- |
| `sparsity` | finite, `[0, 1]` | Nominal dropped-key-block fraction; default `0.8`; `0` disables sparsity |
| `start_step` | integer, `>= 0` | Number of early denoise steps kept dense |
| `end_step` | integer, `>= 0` | Number of final denoise steps kept dense (tail fallback) |
| `precision` | `"bf16"`, `"fp8"`, `"mix"` | Kernel precision mode; default `"bf16"`. Requires MindIE-SD with `sparse_attention(precision=...)` |
| `skip_layers` | selector such as `"0-3,38"` | DiT blocks kept dense |
```bash
@@ -52,6 +54,12 @@ platform raises. It is incompatible with ring sequence parallelism because
the kernel needs the complete key sequence for ranking; use Ulysses sequence
parallelism with `ring_degree=1`.
The `precision` knob requires a MindIE-SD release whose `sparse_attention`
accepts `precision=`; older releases accept it through `**kwargs` but
silently ignore it and run the BF16 path. `RAINFUSION_ATTN` raises at runtime
when a non-`bf16` precision is requested without that support, so pin a
compatible MindIE-SD release when using `precision="mix"` or `"fp8"`.
## Geometry handling
RainFusion handles arbitrary video grids by rearranging video spatially and,
+35 -2
View File
@@ -12,15 +12,17 @@ LoRA adapters are lightweight, model-specific fine-tuning weights that can be ap
## LoRA Adapter Format
### PEFT (Parameter-Efficient Fine-Tuning) format (default)
A typical PEFT-format LoRA adapter directory structure:
```
```text
lora_adapter/
├── adapter_config.json
└── adapter_model.safetensors
```
The `adapter_config.json` file contains metadata about the LoRA adapter, including:
- `r`: LoRA rank
- `lora_alpha`: LoRA alpha scaling factor
- `target_modules`: List of module names to apply LoRA to
@@ -67,7 +69,7 @@ For distilled few-step LoRAs, pass `lora_backend="distill"` together with one or
For Qwen-Image and Wan pipelines, the distill backend calls `pipeline.load_lora_weights(...)` during worker initialization.
| Pipeline | Supported distilled LoRA repo | Notes |
|----------|--------------------------------|-------|
| ---------- | -------------------------------- | ------- |
| `QwenImagePipeline` | `lightx2v/Qwen-Image-2512-Lightning` | Used with Qwen-Image-2512 Lightning-style few-step inference. |
| `Wan22Pipeline` | `lightx2v/Wan2.1-Distill-Loras`, `lightx2v/Wan2.2-Distill-Loras` | Wan2.1 uses one LoRA file. For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. |
| `Wan22I2VPipeline` | `lightx2v/Wan2.2-Distill-Loras` | For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. |
@@ -240,6 +242,37 @@ Notes:
- Output quality and speed depend on the replacement checkpoints and sampling params you choose.
- If you only need to fuse distilled LoRAs into a Wan2.2 checkpoint at load time (without the full LightX2V convert + assemble pipeline), you can instead pass them directly via `--lora-backend distill --lora-path <high>.safetensors <low>.safetensors`. See the [Distill backend](#distill-backend-fuse-distilled-lora-at-init) section above.
## MiniMax-H3 adapter-declared schedules
MiniMax-H3 supports two few-step mechanisms that must not be conflated:
- **Checkpoint-pinned schedule**: a merged release writes `base_schedule` into
`model_index.json`. Requests must pass `num_inference_steps` as the interval
count (for example `4` for `[1.0, 0.7, 0.4, 0.15, 0.0]`).
- **Runtime Turbo LoRA**: LightX2V Turbo artifacts request five sigma points and
enforce `flow_shift=6`, `audio_flow_shift=3`.
- **Runtime native LoRA**: FlashGen-style artifacts declare
`key_format=minimax-h3-native` and embed `base_schedule` in safetensors
metadata. When active, the adapter schedule overrides the base checkpoint and
requests must use the interval-count contract (`num_inference_steps=4`).
Request-mode generation may omit the field and take the count from the adapter
schedule; step execution requires it explicitly, because the step scheduler
derives the total step count from the request before the adapter schedule is
known.
Native artifacts also declare `qkv_layout=grouped`. The H3 loader reorders fused
`qkv_proj` LoRA rows with the same `_reorder_grouped_qkv_to_qkv` helper used for
base-weight loading, then binds the packed Q/K/V slices through the legacy PEFT
manager without modifying `DiffusionLoRAManager`.
Both runtime adapters run with distributed layerwise offload in request-mode
generation, where the manager keeps the LoRA A/B buffers resident on the compute
device while DLO streams the base blocks. MiniMax-H3 step execution still
rejects DLO, so `--step-execution` cannot be combined with
`--enable-distributed-layerwise-offload`. Model-level CPU offload and standard
layerwise offload are rejected in every mode, because the dynamic LoRA tensors
are neither parameters nor registered buffers and therefore do not participate
in those weight lifecycles.
## See Also
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2/LTX-2.3,
@@ -111,6 +111,13 @@ def build_image_to_video_prompt(
return result
def _validate_video_output_type(output_type: str) -> None:
if output_type not in {"image", "video"}:
raise ValueError(
f"Unexpected output type '{output_type}', expected 'video' or legacy 'image' for video generation."
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description=(
@@ -665,10 +672,7 @@ def main():
frames = frames[0] if frames else None
if isinstance(frames, OmniRequestOutput):
if frames.final_output_type != "image":
raise ValueError(
f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
)
_validate_video_output_type(frames.final_output_type)
if frames.multimodal_output and "audio" in frames.multimodal_output:
audio = frames.multimodal_output["audio"]
audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import argparse
import json
@@ -174,6 +174,13 @@ def build_text_to_video_prompt(prompt: str, negative_prompt: str | None) -> dict
return result
def _validate_video_output_type(output_type: str) -> None:
if output_type not in {"image", "video"}:
raise ValueError(
f"Unexpected output type '{output_type}', expected 'video' or legacy 'image' for video generation."
)
def _normalize_float_tensor(tensor: torch.Tensor, source_range: str) -> torch.Tensor:
"""Normalize decoded frames according to the pipeline's output contract."""
if not tensor.is_floating_point():
@@ -620,10 +627,7 @@ def main():
frames = frames[0] if frames else None
if isinstance(frames, OmniRequestOutput):
if frames.final_output_type != "image":
raise ValueError(
f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
)
_validate_video_output_type(frames.final_output_type)
if frames.multimodal_output and "audio" in frames.multimodal_output:
audio = frames.multimodal_output["audio"]
audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
@@ -159,6 +159,33 @@ python examples/online_serving/minicpmo/realtime_duplex_demo.py \
--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
The browser UI serves the page and proxies the same-origin Realtime WebSocket to
@@ -12,7 +12,9 @@ from __future__ import annotations
import argparse
import asyncio
import base64
import io
import json
import math
import sys
import time
from pathlib import Path
@@ -22,8 +24,8 @@ if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
PCM16_BYTES_PER_SAMPLE, # noqa: F401 - compatibility export
PCM16_SAMPLE_RATE, # noqa: F401 - compatibility export
PCM16_BYTES_PER_SAMPLE,
PCM16_SAMPLE_RATE,
RealtimeDuplexClient,
RealtimeEventCollector,
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
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:
@@ -159,14 +165,258 @@ def _event_count_after(
return sum(event.get("type") == event_type for event in events[index + 1 :])
async def run_demo(args: argparse.Namespace) -> dict[str, object]:
input_pcm16 = read_pcm16_wav(Path(args.input_wav))
if not input_pcm16:
raise ValueError("input WAV has no audio")
def _probe_video_frames_and_fps(video_path: Path) -> tuple[int, float]:
"""Frame count and average FPS via PyAV (same approach as Daily-Omni)."""
import av
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)
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(
args.url,
args.model,
@@ -184,10 +434,12 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
timeout_s=args.timeout_s,
)
stream_event_cursor = len(client.events.events)
await client.stream_pcm16(
frames_sent = await client.stream_pcm16(
input_pcm16,
chunk_ms=args.chunk_ms,
realtime=not args.no_realtime_pacing,
video_frames=video_frames,
stacked_video_frames=stacked_video_frames,
)
commit_event_cursor = len(client.events.events)
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,
),
},
"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_chunk_count": len(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("--model", default="openbmb/MiniCPM-o-4_5")
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(
"--ref-audio",
required=True,
@@ -372,7 +669,7 @@ def parse_args() -> argparse.Namespace:
"--temperature",
type=float,
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("--require-audio", action="store_true")
+42 -35
View File
@@ -71,11 +71,11 @@ dev = [
"voxcpm==2.0.3", # for VoxCPM2 path-locating in tests
"torchdiffeq==0.2.5", # for covo_audio Test
"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
# Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo,
# whose top-level stepaudio2 package conflicts with this distribution.
"step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml)
"silero-vad==6.1", # Realtime server-VAD duplex E2E
]
demo = [
@@ -85,6 +85,10 @@ demo = [
"cn2an>=0.5.22", # Chinese number normalization for TTS (e.g. "1次" -> "一次")
]
server-vad = [
"silero-vad==6.1",
]
forced-aligner = [
"qwen-asr>=0.0.6", # faithful Japanese/Korean word segmentation for the TTS forced aligner
]
@@ -244,11 +248,11 @@ addopts = [
"--dist=loadgroup"
]
markers = [
# ci/cd required
"core_model: L1&L2 tests (run in each PR)",
"advanced_model: L3 level tests (run on each merge)",
"full_model: L4 level tests (run nightly)",
"local_model: tests (run locally)",
# ci/cd required — tag with [ci-level] for tests.helpers.mark.get_level_markers().
"core_model: [ci-level] L1&L2 tests (run in each PR)",
"advanced_model: [ci-level] L3 level tests (run on each merge)",
"full_model: [ci-level] L4 level tests (run nightly)",
"local_model: [ci-level] tests (run locally)",
# function module markers
"diffusion: Diffusion model tests",
"omni: Omni model tests",
@@ -258,38 +262,41 @@ markers = [
"sp: Sequence parallelism tests",
"example: Doc example code tests",
# platform markers
"cpu: Tests that run on CPU",
"gpu: Tests that run on GPU (auto-added)",
"cuda: Tests that run on CUDA (auto-added)",
"rocm: Tests that run on AMD/ROCm (auto-added)",
"xpu: Tests that run on XPU (auto-added)",
"npu: Tests that run on NPU/Ascend (auto-added)",
"musa: Tests that run on MUSA/Moore Threads (auto-added)",
# Tag with [hardware-platform] for tests.helpers.mark.get_supported_platforms().
# Add [gpu] on accelerator platforms that should also get pytest.mark.gpu.
# cpu/gpu are pytest marks; hardware_test(res=...) rejects them as keys.
"cpu: [hardware-platform] Tests that run on CPU",
"gpu: [hardware-platform] Tests that run on GPU (auto-added)",
"cuda: [hardware-platform] [gpu] Tests that run on CUDA (auto-added)",
"rocm: [hardware-platform] [gpu] Tests that run on AMD/ROCm (auto-added)",
"xpu: [hardware-platform] [gpu] Tests that run on XPU (auto-added)",
"npu: [hardware-platform] Tests that run on NPU/Ascend (auto-added)",
"musa: [hardware-platform] [gpu] Tests that run on MUSA/Moore Threads (auto-added)",
# specified computation resources marks (auto-added)
# Tag with [hardware-resource] for tests.helpers.mark.get_hardware_mark_list().
"H100: [hardware-resource] Tests that require H100 GPU",
"H800: [hardware-resource] Tests that require H800 GPU",
"H200: [hardware-resource] Tests that require H200 GPU",
"L4: [hardware-resource] Tests that require L4 GPU",
"B200: [hardware-resource] Tests that require B200 GPU",
"MI325: [hardware-resource] Tests that require MI325 GPU (AMD/ROCm)",
"B60: [hardware-resource] Tests that require Intel Arc Pro B60 XPU",
"S5000: [hardware-resource] Tests that require S5000 GPU (Moore Threads/MUSA)",
"A2: [hardware-resource] Tests that require A2 NPU",
"A3: [hardware-resource] Tests that require A3 NPU",
"distributed_cuda: Tests that require multi cards on CUDA platform",
"distributed_rocm: Tests that require multi cards on ROCm platform",
"distributed_xpu: Tests that require multi cards on XPU platform",
"distributed_npu: Tests that require multi cards on NPU platform",
"distributed_musa: Tests that require multi cards on MUSA platform",
"skipif_cuda: Skip if the num of CUDA cards is less than the required",
"skipif_rocm: Skip if the num of ROCm cards is less than the required",
"skipif_xpu: Skip if the num of XPU cards is less than the required",
"skipif_npu: Skip if the num of NPU cards is less than the required",
"skipif_musa: Skip if the num of MUSA cards is less than the required",
# [hardware-resource] + [cuda]/[rocm]/[xpu]/[npu]/[musa] for hardware_test(res=...).
# Do not put [gpu] on SKU lines: that would make gpu a res key.
"H100: [hardware-resource] [cuda] Tests that require H100 GPU",
"H800: [hardware-resource] [cuda] Tests that require H800 GPU",
"H200: [hardware-resource] [cuda] Tests that require H200 GPU",
"L4: [hardware-resource] [cuda] Tests that require L4 GPU",
"B200: [hardware-resource] [cuda] Tests that require B200 GPU",
"MI325: [hardware-resource] [rocm] Tests that require MI325 GPU (AMD/ROCm)",
"B60: [hardware-resource] [xpu] Tests that require Intel Arc Pro B60 XPU",
"S5000: [hardware-resource] [musa] Tests that require S5000 GPU (Moore Threads/MUSA)",
"A2: [hardware-resource] [npu] Tests that require A2 NPU",
"A3: [hardware-resource] [npu] Tests that require A3 NPU",
# Tag with [hardware-cards] for tests.helpers.mark.get_supported_card_counts().
"cards_1: [hardware-cards] Tests that require 1 accelerator card",
"cards_2: [hardware-cards] Tests that require 2 accelerator cards",
"cards_3: [hardware-cards] Tests that require 3 accelerator cards",
"cards_4: [hardware-cards] Tests that require 4 accelerator cards",
"cards_5: [hardware-cards] Tests that require 5 accelerator cards",
"cards_6: [hardware-cards] Tests that require 6 accelerator cards",
"cards_7: [hardware-cards] Tests that require 7 accelerator cards",
"cards_8: [hardware-cards] Tests that require 8 accelerator cards",
# more detailed markers
"sp: Sequence parallelism tests (multi-GPU)",
"slow: Slow tests (may skip in quick CI)",
"slow: [ci-level] Slow tests (may skip in quick CI)",
"benchmark: Benchmark tests",
"xdist: xdist-parallelizable tests",
]
+47 -5
View File
@@ -10,14 +10,22 @@ checkpoint for T2V and first-frame I2V. Every output includes synchronized
48 kHz stereo audio.
The raw [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5)
repository is not directly loadable with `--model`. It supplies the official
upsampler and LoRA sidecars used by the Full/SFT two-stage pipeline, so accept
both model licenses and authenticate before first use.
repository is not directly loadable with `--model`. It supplies the canonical
Diffusion VAE checkpoint and the official upsampler and LoRA sidecars used by
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
| Pipeline | Mode | Default size | Schedule |
|---|---|---:|---:|
| --- | --- | ---: | ---: |
| `LTX2Pipeline` | Full/SFT one-stage | 960x544 | 30 steps |
| `LTX2TwoStagePipeline` | Full/SFT two-stage | 1920x1088 | 30 + 3 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
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
```bash
@@ -42,7 +84,7 @@ Install matching vLLM and vLLM-Omni versions, and ensure `ffmpeg` and
## Hardware
| GPU | Status | Recommended scope |
|---|---|---|
| --- | --- | --- |
| NVIDIA B300 | 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 |
@@ -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.
+57
View File
@@ -860,6 +860,63 @@ fixed HBM usage. Model-level and standard layerwise offload remain unsupported.
The five requested sigma points produce the four denoiser evaluations expected
by the Turbo artifact.
### FlashGen native LoRA
The FlashGen 4-step T2VA artifact uses the native MiniMax-H3 module layout and
declares its distilled sigma schedule in safetensors metadata. It is published on
[ModelScope](https://modelscope.cn/models/FlashGen/Minimax-H3-4step-lora-flashgen):
```text
FlashGen/Minimax-H3-4step-lora-flashgen/minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors
```
Download only that file:
```bash
python -m pip install modelscope
export FLASHGEN_DIR=/path/to/minimax-h3-flashgen-lora
export FLASHGEN_FILE=minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors
modelscope download FlashGen/Minimax-H3-4step-lora-flashgen \
--local_dir "${FLASHGEN_DIR}" \
--include "${FLASHGEN_FILE}"
export FLASHGEN_LORA="${FLASHGEN_DIR}/${FLASHGEN_FILE}"
```
Start from a non-offloaded or DLO FL2VA server command and add
`--task-type fl2va --lora-backend peft --lora-path "${FLASHGEN_LORA}"`.
Each request must use T2VA and the distilled interval-count contract:
```bash
-F 'num_inference_steps=4' \
-F 'extra_params={"task":"t2va","duration":5.2}' \
-F "lora={\"name\":\"h3-flashgen-v1.0\",\"path\":\"${FLASHGEN_LORA}\",\"scale\":1.0}"
```
This path rejects Ref2VA and checkpoints that already pin `base_schedule` in
`model_index.json`. The adapter metadata carries
`base_schedule=1.0,0.7,0.4,0.15,0.0`, so `num_inference_steps=4` means four
denoiser evaluations, not five sigma points. Request-mode generation may omit
the field and take the count from the adapter schedule; `--step-execution`
requires it explicitly, because the step scheduler reads the total step count
off the request at admission, before the adapter schedule is known.
DLO is supported in request-mode generation on the same terms as the Turbo
adapter: the request-switchable LoRA A/B buffers stay resident on the
accelerator while DLO streams only the base blocks, so budget for that
additional fixed HBM usage. The native artifact is rank 64 over 259 target
modules, and its packed `qkv_proj` and `fc1` layers reuse the full-input A
tensor per slice while B carries slice-local output rows, so the resident
footprint exceeds the on-disk payload; measure it for your parallel layout
rather than assuming the checkpoint size. Pure Ulysses replicates the adapter
on every rank, while DiT tensor parallelism shards the B buffers. Model-level
and standard layerwise offload remain unsupported, and `--step-execution`
cannot be combined with `--enable-distributed-layerwise-offload`.
To validate a deployment, post the same fixed-seed T2VA request twice with the
adapter and twice without it, then compare the four output digests. The adapter
is bound and deterministic when each pair matches internally and the two pairs
differ from each other.
## Key parameters
| Parameter | Recommended value | Notes |
+14 -5
View File
@@ -18,11 +18,11 @@ examples in this repository.
Qwen3-TTS supports three task types, each backed by a dedicated model checkpoint:
| Task Type | Model | Description |
| ------------- | ---------------------------------------- | ------------------------------------------------------------ |
| `CustomVoice` | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speaker voices with optional style/emotion control |
| `VoiceDesign` | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Generate speech from a natural language voice description |
| `Base` | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from reference audio + transcript |
| Task Type | Model | Description |
| ------------- | ---------------------------------------- | ------------------------------------------------------------- |
| `CustomVoice` | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speaker voices with optional style/emotion control |
| `VoiceDesign` | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Generate speech from a natural language voice description |
| `Base` | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from reference audio + transcript |
Smaller 0.6B variants are also available for `CustomVoice` and `Base`.
@@ -166,6 +166,15 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-ty
- Key flags: `--omni` is required. `--deploy-config` points to the bundled two-stage pipeline config.
- Async chunking: Enabled by default in `qwen3_tts.yaml` for streaming-friendly first-audio latency. Raw audio streaming requires `stream=true`, `stream_format="audio"`, and `response_format="pcm"`.
- Task/model matching: Each task type requires its matching model checkpoint. Using a CustomVoice model for a Base (voice clone) request will fail.
- Base codec termination: Base requests without an explicit `max_new_tokens`
use a text-scaled safety ceiling (at least 192 codec frames and no more than
the configured model limit). If the Talker reaches that ceiling without
codec EOS, non-streaming serving discards the incomplete audio and retries
once with a fresh seed; an explicit `seed` or `max_new_tokens` disables the
retry. SSE and WebSocket clients receive structured errors and must discard
previously emitted audio when the error contains `"action":"discard"`;
raw PCM streams terminate with a connection error after any partial bytes
already sent.
- Known limitations: The server serves one model variant at a time. To switch task types (e.g., CustomVoice to Base), restart the server with the corresponding model.
## Hardware Support
+2
View File
@@ -13,6 +13,8 @@ transformers >= 5.10.1, < 5.15
av>=14.0.0
omegaconf>=2.3.0
diffusers==0.40.0
# Hub kernels, including LTX-2.5 NATTEN
kernels==0.14.1
safetensors>=0.8.0
accelerate==1.12.0
soundfile>=0.13.1
+66 -4
View File
@@ -279,17 +279,79 @@ def test_unmeasured_tpot_stays_missing_after_tokenizer_fallback():
assert metrics.request_goodput == 0.0
def test_measured_zero_itl_is_not_treated_as_missing():
def test_zero_itl_does_not_create_zero_tpot():
output = _make_output(100, output_tokens=3)
output.itl = [0.0, 0.0]
output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
metrics = _calculate_test_metrics([output], {"tpot": 1.0})
assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (1, 2)
assert metrics.mean_tpot_ms == 0.0
assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (0, 2)
assert math.isnan(metrics.mean_tpot_ms)
assert metrics.mean_itl_ms == 0.0
assert metrics.request_goodput == 1.0
assert metrics.request_goodput == 0.0
def test_stage_output_tokens_without_client_timing_omit_tpot(capsys):
output = _make_output(100, output_tokens=5)
output.itl = []
output.text_latency = 0.0
output.ttft = 0.1
metrics, _ = calculate_metrics(
input_requests=[],
outputs=[output],
dur_s=1.0,
tokenizer=None,
selected_percentiles=[50.0, 99.0],
goodput_config_dict={},
task_type=TaskType.GENERATION,
selected_percentile_metrics=["tpot"],
max_concurrency=None,
request_rate=float("inf"),
benchmark_duration=1.0,
)
printed = capsys.readouterr().out
assert metrics.num_tpot_samples == 0
assert math.isnan(metrics.mean_tpot_ms)
assert "Time per Output Token" not in printed
assert "Mean TPOT" not in printed
def test_consistent_client_latency_can_supply_tpot_fallback():
output = _make_output(100, output_tokens=10)
output.itl = []
output.ttft = 0.1
output.text_latency = 1.0
metrics = _calculate_test_metrics([output])
assert metrics.num_tpot_samples == 1
assert metrics.mean_tpot_ms == pytest.approx(100.0)
def test_single_token_responses_do_not_report_zero_tpot(capsys):
output = _make_output(100, output_tokens=1)
metrics, _ = calculate_metrics(
input_requests=[],
outputs=[output],
dur_s=1.0,
tokenizer=None,
selected_percentiles=[50.0],
goodput_config_dict={},
task_type=TaskType.GENERATION,
selected_percentile_metrics=["tpot"],
max_concurrency=None,
request_rate=float("inf"),
benchmark_duration=1.0,
)
printed = capsys.readouterr().out
assert metrics.num_tpot_samples == 0
assert math.isnan(metrics.mean_tpot_ms)
assert "Time per Output Token" not in printed
def test_duplex_goodput_does_not_pair_measurements_from_different_requests():
+73 -2
View File
@@ -1,5 +1,5 @@
# 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 patch.py
@@ -17,6 +17,7 @@ from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput
from vllm_omni.benchmarks.patch.patch import (
MixRequestFuncOutput,
_apply_stage0_token_timings,
async_request_openai_chat_omni_completions,
async_request_openai_realtime_duplex,
)
@@ -102,7 +103,24 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
received_at_s=now + 0.03,
)
self.events.add(
{"type": "response.done", "response": {"id": response_id}},
{
"type": "response.done",
"response": {
"id": response_id,
"metadata": {
"vllm_omni": {
"stage_metrics": {
"0": {
"num_tokens_out": 3,
"vllm_ttft_ms": 20.0,
"vllm_tpot_ms": 10.0,
"vllm_itls_ms": [10.0, 10.0],
}
}
}
},
},
},
received_at_s=now + 0.04,
)
@@ -143,6 +161,7 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
assert client.configure_kwargs["native_duplex"] is False
assert client.configure_kwargs["extra_body"] == {
"ref_audio": "data:audio/wav;base64,AAAA",
"return_stage_metrics": True,
}
assert client.sent == [
event
@@ -167,6 +186,10 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
assert output.audio_ttfp > output.ttft
assert output.audio_rtf > 0
assert output.latency > 0
assert output.output_tokens == 12
assert output.itl == [0.01, 0.01] * 4
assert output.text_latency == pytest.approx(output.ttft + 0.08)
assert output.tpot_measured is True
assert output.tts_turn_pcm_bytes == [b"\x00\x00" * 2400] * 4
assert output.tts_output_pcm_bytes == b"\x00\x00" * 9600
session_id = output.duplex_request_metrics[0]["session_id"]
@@ -198,6 +221,54 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
}
def test_stage0_token_timings_use_weighted_tpot_when_itls_are_incomplete():
output = MixRequestFuncOutput(ttft=0.1)
measured = _apply_stage0_token_timings(
output,
[
{"output_token_count": 3, "itls_ms": [], "tpot_ms": 10.0},
{"output_token_count": 2, "itls_ms": [], "tpot_ms": 40.0},
],
expected_output_tokens=5,
)
assert measured is True
assert output.itl == []
assert output.tpot_measured is True
assert (output.text_latency - output.ttft) / 4 == pytest.approx(0.02)
def test_stage0_zero_itls_fall_back_to_positive_tpot():
output = MixRequestFuncOutput(ttft=0.1)
measured = _apply_stage0_token_timings(
output,
[{"output_token_count": 3, "itls_ms": [0.0, 0.0], "tpot_ms": 25.0}],
expected_output_tokens=3,
)
assert measured is True
assert output.itl == []
assert output.tpot_measured is True
assert (output.text_latency - output.ttft) / 2 == pytest.approx(0.025)
def test_stage0_token_count_without_timing_is_not_measured():
output = MixRequestFuncOutput(ttft=0.1)
measured = _apply_stage0_token_timings(
output,
[{"output_token_count": 3, "itls_ms": [], "tpot_ms": None}],
expected_output_tokens=3,
)
assert measured is False
assert output.itl == []
assert output.text_latency == output.ttft
assert output.tpot_measured is False
def create_sse_chunk(data_dict):
"""Helper to create SSE formatted chunk"""
return f"data: {json.dumps(data_dict)}\n\n".encode()
@@ -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.
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:
_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)
sys.modules[_MODULE_NAME] = _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):
counter.append(1)
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
monkeypatch.setattr(dataset, "_extract_minicpm_frame_audio_segments", fake_extract)
+6 -2
View File
@@ -585,7 +585,7 @@ async def test_adapter_requires_and_forwards_exact_prepared_payload(tmp_path: Pa
@pytest.mark.asyncio
@pytest.mark.parametrize(
("request_metrics", "expected_itl", "expected_text_latency"),
("request_metrics", "expected_itl", "expected_text_latency", "expected_tpot_measured"),
[
(
[
@@ -594,6 +594,7 @@ async def test_adapter_requires_and_forwards_exact_prepared_payload(tmp_path: Pa
],
[0.01, 0.02, 0.03],
0.16,
True,
),
(
[
@@ -602,6 +603,7 @@ async def test_adapter_requires_and_forwards_exact_prepared_payload(tmp_path: Pa
],
[],
0.18,
True,
),
(
[
@@ -610,6 +612,7 @@ async def test_adapter_requires_and_forwards_exact_prepared_payload(tmp_path: Pa
],
[],
0.1,
False,
),
],
)
@@ -619,6 +622,7 @@ async def test_adapter_reports_exact_or_weighted_token_timing(
request_metrics: list[dict[str, object]],
expected_itl: list[float],
expected_text_latency: float,
expected_tpot_measured: bool,
):
case, options = _case(tmp_path), _session_options(tmp_path)
request = _request(case, options)
@@ -643,7 +647,7 @@ async def test_adapter_reports_exact_or_weighted_token_timing(
assert output.success
assert output.itl == expected_itl
assert output.text_latency == pytest.approx(expected_text_latency)
assert output.tpot_measured is True
assert output.tpot_measured is expected_tpot_measured
def test_batch_finalization_publishes_only_measured_results(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+142 -1
View File
@@ -34,6 +34,7 @@ from vllm_omni.config.stage_config import (
build_stage_runtime_overrides,
load_deploy_config,
merge_pipeline_deploy,
normalize_pipeline_cli_overrides,
pipeline_cfg_resolver,
)
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_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):
"""Test diffusion config calculates devices from parallel_config."""
@@ -1144,7 +1159,7 @@ class TestDeployConfigLoading:
stages = merge_pipeline_deploy(pipeline, deploy)
assert deploy.session_mode == "duplex"
assert deploy.active_stream_window == 1
assert deploy.active_stream_window == max_sessions
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"]
@@ -1324,6 +1339,13 @@ stages:
("step_audio2_ci.yaml", "step_audio_2_asr", 1, "text", "step_audio_2_asr"),
("hunyuan_video_15.yaml", "hunyuan_video_15", 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(
@@ -1344,9 +1366,91 @@ stages:
assert len(stages) == stage_count
assert stages[-1].final_output is True
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:
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(
("config_json", "model_index", "expected_pipeline"),
[
@@ -1382,6 +1486,29 @@ stages:
assert pipeline is not None
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"])
def test_step_audio2_deploy_configs_fit_two_gpus(self, deploy_name):
deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
@@ -2272,6 +2399,20 @@ class TestPlatformOverrides:
assert rocm.stages[0].enforce_eager is None
assert rocm.stages[1].enforce_eager is True
def test_higgs_audio_v3_rocm_uses_triton_attention(self):
deploy_path = Path(get_deploy_config_path("higgs_multimodal_qwen3.yaml"))
base = load_deploy_config(deploy_path)
assert base.stages[0].engine_extras["attention_backend"] == "FLASHINFER"
rocm = _apply_platform_overrides(base, platform="rocm")
assert rocm.stages[0].engine_extras["attention_backend"] == "TRITON_ATTN"
pipeline = resolve_pipeline_config("higgs_multimodal_qwen3")
assert isinstance(pipeline, PipelineConfig)
stages = merge_pipeline_deploy(pipeline, rocm)
assert stages[0].yaml_engine_args["attention_backend"] == "TRITON_ATTN"
def test_qwen3_omni_cuda_uses_thinker_rotary_custom_op(self):
deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
pipeline = resolve_pipeline_config(
+64 -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:
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.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:
@@ -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))
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))
@@ -132,6 +146,7 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod
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.session_mode == engine_args.get("session_mode", "turn")
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.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
@@ -579,6 +594,7 @@ def test_sub_config_fields_match_structured_scopes():
"interleave_mm_strings",
"media_io_kwargs",
"active_stream_window",
"session_mode",
"duplex_max_sessions",
"enable_sleep_mode",
"default_sampling_params",
@@ -637,10 +653,11 @@ def test_sub_config_fields_match_structured_scopes():
"ring_degree",
"allgather_degree",
"ulysses_mode",
"ulysses_a2a_permute",
"cfg_parallel_size",
"vae_patch_parallel_size",
"text_encoder_tp_size",
"vae_parallel_mode",
"text_encoder_tp_size",
"use_hsdp",
"mask_sp_padding",
"hsdp_shard_size",
@@ -842,6 +859,23 @@ def test_from_pipeline_config_derives_sequence_parallel_size_from_allgather_degr
assert stage.parallel_config.world_size == 2
def test_from_pipeline_config_preserves_deployed_ulysses_a2a_permute() -> None:
pipeline = _resolve_pipeline_or_skip("dreamzero")
deploy = DeployConfig(
stages=[
StageDeployConfig(
stage_id=0,
ulysses_a2a_permute=True,
)
]
)
stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0)
assert isinstance(stage, VllmOmniDiffusionStageConfig)
assert stage.parallel_config.ulysses_a2a_permute is True
def test_diffusion_parallel_config_accepts_four_way_guidance_parallelism():
cfg = OmniStageDiffusionParallelConfig(cfg_parallel_size=4)
@@ -936,6 +970,35 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
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():
with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
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."""
with pytest.raises(ValueError):
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"
+2 -2
View File
@@ -25,9 +25,9 @@ pytest_plugins = (
"tests.helpers.fixtures.clean",
"tests.helpers.fixtures.log",
"tests.helpers.fixtures.media",
"tests.helpers.fixtures.run_args",
"tests.helpers.fixtures.pytest_collection",
"tests.helpers.fixtures.pytest_run_args",
"tests.helpers.fixtures.runtime",
"tests.helpers.fixtures.xdist",
)
@@ -0,0 +1,89 @@
"""Tests for ``OmniARScheduler._drop_aborted_queued_requests``.
``schedule()`` sweeps ``FINISHED_ABORTED`` requests out of the queues before
handing control to upstream, because upstream ``Scheduler.schedule()`` raises
``RuntimeError: Invalid request status: FINISHED_ABORTED`` for any request it
admits in a finished state -- which kills the stage's engine core, not just
the request.
Regression for the sweep missing ``skipped_waiting``, the third queue
upstream admits from. An aborted duplex session parked there was re-selected
by ``_select_waiting_queue_for_scheduling`` on a later tick and crashed the
stage.
"""
from __future__ import annotations
import pytest
from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
from vllm.v1.request import RequestStatus
import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
_POLICIES = [
pytest.param(SchedulingPolicy.FCFS, id="fcfs"),
pytest.param(SchedulingPolicy.PRIORITY, id="priority"),
]
class _StubRequest:
"""Minimal ``Request`` stub with the surface the sweep exercises.
Carries the ``(priority, arrival_time)`` ordering ``PriorityRequestQueue``
relies on so the same stub works under both scheduling policies.
"""
def __init__(self, request_id: str, status: RequestStatus, arrival_time: float = 0.0) -> None:
self.request_id = request_id
self.status = status
self.priority = 0
self.arrival_time = arrival_time
def __lt__(self, other: _StubRequest) -> bool:
return (self.priority, self.arrival_time) < (other.priority, other.arrival_time)
def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), skipped_waiting=(), running=()):
"""Bare scheduler with just the surface the sweep reads."""
scheduler = OmniARScheduler.__new__(OmniARScheduler)
scheduler.waiting = create_request_queue(policy)
for req in waiting:
scheduler.waiting.add_request(req)
scheduler.skipped_waiting = create_request_queue(policy)
for req in skipped_waiting:
scheduler.skipped_waiting.add_request(req)
scheduler.running = list(running)
return scheduler
@pytest.mark.parametrize("policy", _POLICIES)
def test_schedule_sweeps_skipped_waiting_before_upstream_selection(
monkeypatch: pytest.MonkeyPatch,
policy: SchedulingPolicy,
) -> None:
"""The public schedule path must remove aborted skipped requests first."""
aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED)
scheduler = _make_scheduler(policy, skipped_waiting=[aborted])
scheduler.chunk_transfer_adapter = None
scheduler.input_coordinator = None
scheduler.num_waiting_for_streaming_input = 0
scheduler._process_pending_omni_inputs = lambda model_mode: None
scheduler._resync_streaming_input_counter = lambda: None
scheduler._should_defer_waiting_admission = lambda: False
scheduler._restore_omni_wait_queues = lambda: None
scheduler._postprocess_omni_schedule_output = lambda *args, **kwargs: None
scheduler.get_finished_requests_needing_kv_transfer = lambda: {}
scheduler._wrap_omni_scheduler_output = lambda output, **kwargs: output
sentinel = object()
def upstream_schedule(_self, _throttle_prefills: bool = False):
assert list(scheduler.skipped_waiting) == []
return sentinel
monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule)
assert OmniARScheduler.schedule(scheduler) is sentinel
@@ -142,13 +142,16 @@ def test_exact_drain_delivers_new_segment_frame() -> None:
session.num_computed_tokens = 6
session.num_output_placeholders = 1
session.num_in_flight_tokens = 1
session.async_tokens_to_discard = 1
_replace_streaming_session(session)
assert session.num_stale_output_tokens == 1
assert session.async_tokens_to_discard == 1
sched = _make_drain_sched(session)
assert _run_step(sched, session, num_scheduled=1, token=42) is False # late frame dropped
assert session.num_stale_output_tokens == 0
assert session.async_tokens_to_discard == 0
assert _run_step(sched, session, num_scheduled=1, token=43) is True # new segment survives
@@ -18,17 +18,19 @@ 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, StreamingUpdate
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
# isort: on
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._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.log_stats = False
sched.chunk_transfer_adapter = None
@@ -185,6 +187,7 @@ def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> Non
def replace_with_next_segment(request: Request) -> bool:
request.num_computed_tokens = 0
request._omni_segment_generation = 1
request.status = RequestStatus.WAITING
return False
@@ -201,6 +204,7 @@ def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> Non
True,
new_token_ids=[42],
confirmed_num_computed_tokens=26,
segment_generation=0,
)
@@ -286,6 +290,140 @@ def test_queued_streaming_update_on_async_stop_fences_in_flight_once() -> None:
assert session.num_stale_output_tokens == session.num_in_flight_tokens == 1
def test_stale_async_frame_is_dropped_before_output_processing() -> None:
session = _make_request()
session.status = RequestStatus.RUNNING
session.num_in_flight_tokens = 2
session.num_computed_tokens = session.num_tokens + 1
session.num_output_placeholders = 1
session.async_tokens_to_discard = 1
session.sampling_params = SimpleNamespace(num_logprobs=1)
num_computed_tokens = session.num_computed_tokens
num_output_placeholders = session.num_output_placeholders
sched = MagicMock()
sched.requests = {session.request_id: session}
sched.perf_metrics = None
sched.structured_output_manager.should_advance.return_value = False
def discard_stale_output(request: Request, token_ids: list[int]) -> tuple[list[int], bool]:
request.async_tokens_to_discard = 0
return token_ids, False
sched._update_request_with_output.side_effect = discard_stale_output
sched._process_kv_transfer_trigger.return_value = False
sched.chunk_transfer_adapter = MagicMock()
sched.running = [session]
sched.waiting_for_transfer_free = set()
sched.transfer_triggered_requests = set()
sched.active_kv_transfers = set()
sched.pending_stop_after_extraction = set()
sched.connector = 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.make_stats.return_value = None
scheduler_output = MagicMock(spec=SchedulerOutput)
scheduler_output.num_scheduled_tokens = {session.request_id: 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.multimodal_outputs = None
model_runner_output.inter_stage_outputs = [{"hidden": object()}]
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 = {session.request_id: 0}
model_runner_output.routed_experts = None
OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
assert session.async_tokens_to_discard == 0
assert session.status == RequestStatus.RUNNING
assert session.num_computed_tokens == num_computed_tokens
assert session.num_output_placeholders == num_output_placeholders
sched.chunk_transfer_adapter.save_async.assert_not_called()
session.sampling_params = SamplingParams(max_tokens=8)
next_payload = {"hidden": object()}
sched._update_request_with_output.side_effect = None
sched._update_request_with_output.return_value = ([43], False)
model_runner_output.sampled_token_ids = [[43]]
model_runner_output.inter_stage_outputs = [next_payload]
OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
sched.chunk_transfer_adapter.save_async.assert_called_once_with(
next_payload,
session,
False,
new_token_ids=[43],
confirmed_num_computed_tokens=None,
)
def test_legacy_stale_async_marker_bypasses_real_placeholder_accounting() -> None:
request = _make_request()
request.status = RequestStatus.RUNNING
request.num_in_flight_tokens = 1
request.num_output_placeholders = 0
request.async_tokens_to_discard = 1
request.num_stale_output_tokens = 0
sched = OmniARAsyncScheduler.__new__(OmniARAsyncScheduler)
sched.requests = {request.request_id: request}
sched.perf_metrics = None
sched.chunk_transfer_adapter = None
sched.connector = None
sched.ec_connector = None
sched.waiting_for_transfer_free = set()
sched.transfer_triggered_requests = set()
sched.active_kv_transfers = set()
sched.pending_stop_after_extraction = set()
sched.finished_req_ids_dict = {}
sched._new_prompt_len_snapshot = {}
sched.kv_cache_manager = MagicMock()
sched.kv_cache_manager.take_events.return_value = None
sched._remove_stopped_requests_from_queues = MagicMock()
sched._handle_failed_kv_load_outputs = MagicMock(return_value=[])
sched._cleanup_kv_tracking = MagicMock()
sched._aggregate_kv_connector_stats = MagicMock(return_value=None)
sched._publish_kv_cache_events = MagicMock()
sched._attach_finished_request_sets = MagicMock()
sched._attach_scheduler_stats = MagicMock()
sched._capture_omni_connector_output = MagicMock()
scheduler_output = MagicMock(spec=SchedulerOutput)
scheduler_output.num_scheduled_tokens = {request.request_id: 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.multimodal_outputs = None
model_runner_output.inter_stage_outputs = 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
OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
assert request.async_tokens_to_discard == 0
assert request.num_output_placeholders == 0
assert request._output_token_ids == []
def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
sched = _make_scheduler(stage_id=0)
session = _make_request()
@@ -446,11 +584,15 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
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(segment_finished_requests={session.request_id})
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})
@@ -460,3 +602,37 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
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()
@@ -76,6 +76,9 @@ def _make_scheduler(
else:
hf_config.voxcpm2_runtime_config = runtime_config
sched.vllm_config = SimpleNamespace(model_config=model_config)
# Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the
# abort sweep at the top of ``schedule()`` reads it.
sched.skipped_waiting = _MockQueue()
return sched
@@ -166,3 +166,15 @@ def test_fast_path_schedule_defer_disabled_leaves_fence() -> None:
OmniGenerationScheduler.schedule(sched)
assert sched.sched_step_seq == 0
def test_generation_schedule_sweeps_and_resyncs_around_pending_inputs() -> None:
sched, request = _make_fast_path_sched(defer_block_free=False)
calls: list[str] = []
sched._drop_aborted_queued_requests.side_effect = lambda: calls.append("sweep")
sched._process_pending_omni_inputs.side_effect = lambda model_mode: calls.append(model_mode)
sched._resync_streaming_input_counter.side_effect = lambda: calls.append("resync")
OmniGenerationScheduler.schedule(sched)
assert calls == ["sweep", "generation", "sweep", "resync"]
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit coverage for the two stage-input timeout hooks.
``_process_pending_input_timeouts`` covers the full-payload transport;
@@ -253,6 +253,38 @@ def test_log_failed_chunk_sends_noop_without_adapter():
scheduler._log_failed_chunk_sends()
class _FakeReceiveFailureAdapter:
def __init__(self, failures):
self._failures = dict(failures)
def collect_failed_receive_request_ids(self):
failures, self._failures = self._failures, {}
return failures
class _FakeReceiveFailureScheduler(OmniSchedulerMixin):
def __init__(self, requests, adapter):
self.requests = requests
self.chunk_transfer_adapter = adapter
self.finish_calls = []
def finish_requests(self, req_ids, status):
self.finish_calls.append((set(req_ids), status))
def test_process_chunk_receive_failures_delegates_to_finish_requests():
request_id = "invalid-chunk"
adapter = _FakeReceiveFailureAdapter({request_id: "prompt exceeds limit"})
scheduler = _FakeReceiveFailureScheduler({request_id: SimpleNamespace(request_id=request_id)}, adapter)
scheduler._process_chunk_receive_failures()
assert len(scheduler.finish_calls) == 1
finished_ids, status = scheduler.finish_calls[0]
assert finished_ids == {request_id}
assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
# --- R1.4: VLLM_OMNI_INPUT_WAIT_TIMEOUT_S is parsed at import time, so each of
# these reloads the module under a different environment. One knob now gates the
# deadline on both transports, which is why a silent 0 matters more than it did
@@ -0,0 +1,168 @@
"""Tests for the ``num_waiting_for_streaming_input`` resync placed in
``OmniARScheduler.finish_requests`` / ``OmniGenerationScheduler.finish_requests``
and in ``OmniARScheduler.schedule``.
Upstream keeps that counter incrementally (``+1`` when a resumable request
parks in a waiting queue as ``WAITING_FOR_STREAMING_REQ``, ``-1`` when the
next update arrives or the request is finished) and
``Scheduler.get_num_unfinished_requests`` *subtracts* it from the waiting
queues. Omni rewrites ``request.status`` outside both hooks -- the
chunk-transfer adapter's park/restore and
``_realign_request_status_to_queues`` -- so a counted request can leave
``WAITING_FOR_STREAMING_REQ`` without the matching decrement and inflate the
counter for the lifetime of the process.
An inflated counter does not degrade throughput, it hangs the stage:
``EngineCore.has_work()`` reads false while a live request sits in
``waiting``, so the engine blocks in ``input_queue.get()`` and never calls
``schedule()`` again. Regression for the Qwen3-Omni duplex symptom where the
server answered the first conversation after a boot and was silent from the
second on while still reporting healthy.
"""
from __future__ import annotations
import pytest
from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
from vllm.v1.request import RequestStatus
import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _StubRequest:
def __init__(self, request_id: str, status: RequestStatus) -> None:
self.request_id = request_id
self.status = status
def is_finished(self) -> bool:
return RequestStatus.is_finished(self.status)
_SCHEDULER_PARAMS = [
pytest.param(OmniARScheduler, ar_sched_mod, id="ar"),
pytest.param(OmniGenerationScheduler, gen_sched_mod, id="generation"),
]
def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipped=None):
scheduler = scheduler_cls.__new__(scheduler_cls)
scheduler.chunk_transfer_adapter = None
scheduler.input_coordinator = None
scheduler.requests = requests
scheduler.running = running
scheduler.waiting = waiting
scheduler.skipped_waiting = skipped if skipped is not None else []
scheduler.num_waiting_for_streaming_input = counter
return scheduler
@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
def test_finish_requests_clears_counter_left_by_a_stomped_status(
monkeypatch: pytest.MonkeyPatch,
scheduler_cls,
scheduler_mod,
) -> None:
"""The abort case that wedged the talker.
The request was counted while parked for streaming input, then had its
status rewritten to ``RUNNING`` by omni's queue surgery, so upstream's
``finish_requests`` took the running branch and never decremented. With
no parked requests left, the counter must land at 0.
"""
aborted = _StubRequest("req-aborted", RequestStatus.RUNNING)
scheduler = _make_scheduler(
scheduler_cls,
requests={}, # upstream already dropped it
running=[],
waiting=[],
counter=1,
)
monkeypatch.setattr(
scheduler_mod.VLLMScheduler,
"finish_requests",
lambda self, request_ids, finished_status: [aborted],
)
scheduler_cls.finish_requests(scheduler, [aborted.request_id], RequestStatus.FINISHED_ABORTED)
assert scheduler.num_waiting_for_streaming_input == 0
def test_resync_counts_parked_requests_in_skipped_waiting() -> None:
"""``get_num_unfinished_requests`` sums ``waiting`` and ``skipped_waiting``,
so the resync has to look at both or it under-counts and reintroduces the
spin the counter exists to prevent."""
parked = _StubRequest("req-skipped", RequestStatus.WAITING_FOR_STREAMING_REQ)
scheduler = _make_scheduler(
OmniARScheduler,
requests={parked.request_id: parked},
running=[],
waiting=[],
counter=0,
skipped=[parked],
)
scheduler._resync_streaming_input_counter()
assert scheduler.num_waiting_for_streaming_input == 1
def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.MonkeyPatch) -> None:
"""Pins placement in ``schedule``: the chunk-transfer adapter rewrites
status during ``process_pending_chunks``, so the resync has to run after
that and before upstream schedules, while ``self.waiting`` is still the
real queue."""
parked = _StubRequest("req-resumed", RequestStatus.RUNNING)
scheduler = OmniARScheduler.__new__(OmniARScheduler)
scheduler.chunk_transfer_adapter = None
scheduler.input_coordinator = None
scheduler.requests = {parked.request_id: parked}
scheduler.running = [parked]
scheduler.waiting = []
scheduler.skipped_waiting = []
scheduler.num_waiting_for_streaming_input = 1 # leaked by the status rewrite
scheduler.policy = "fcfs"
seen: dict[str, int] = {}
monkeypatch.setattr(
ar_sched_mod.OmniSchedulerMixin, "_consume_pending_connector_output", lambda self, model_mode: None
)
monkeypatch.setattr(ar_sched_mod.OmniSchedulerMixin, "_process_pending_input_timeouts", lambda self: None)
monkeypatch.setattr(OmniARScheduler, "_should_defer_waiting_admission", lambda self: False)
def fake_schedule(self, throttle_prefills=False):
seen["counter"] = self.num_waiting_for_streaming_input
raise _StopScheduleError
class _StopScheduleError(Exception):
pass
monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", fake_schedule)
with pytest.raises(_StopScheduleError):
OmniARScheduler.schedule(scheduler)
assert seen["counter"] == 0
def test_generation_scheduler_sweep_removes_aborted_requests_from_all_queues() -> None:
aborted = _StubRequest("req-generation-aborted", RequestStatus.FINISHED_ABORTED)
scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
scheduler.waiting.add_request(aborted)
scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
scheduler.skipped_waiting.add_request(aborted)
scheduler.running = [aborted]
scheduler._drop_aborted_queued_requests()
assert list(scheduler.waiting) == []
assert list(scheduler.skipped_waiting) == []
assert scheduler.running == []
+16
View File
@@ -1,4 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import json
import math
import os
import threading
from collections.abc import Callable
@@ -159,6 +163,18 @@ def benchmark_params(request):
def assert_result(result, params, num_prompt) -> None:
assert result["completed"] == num_prompt, "Request failures exist"
baseline = params.get("baseline")
hardware = result.get("Hardware")
hardware_baseline = baseline.get(hardware) if isinstance(baseline, dict) and isinstance(hardware, str) else None
if isinstance(hardware_baseline, dict) and "mean_tpot_ms" in hardware_baseline:
num_tpot_samples = result.get("num_tpot_samples")
mean_tpot_ms = result.get("mean_tpot_ms")
assert isinstance(num_tpot_samples, int) and not isinstance(num_tpot_samples, bool) and num_tpot_samples > 0, (
"TPOT baseline is configured, but no measurable TPOT samples were produced"
)
assert (
isinstance(mean_tpot_ms, int | float) and not isinstance(mean_tpot_ms, bool) and math.isfinite(mean_tpot_ms)
), "TPOT baseline is configured, but mean_tpot_ms is not finite"
expected_audio_turns = params.get("expected_duplex_audio_turns_per_session")
if expected_audio_turns is not None:
session_metrics = result.get("duplex_session_metrics")
@@ -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
}
}
]
}
]
@@ -43,7 +43,7 @@
"extra_body": {
"save_duplex_request_metrics": true
},
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration",
"percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration",
"baseline": {
"H100": {
"request_throughput": [
@@ -46,9 +46,6 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
"mean_e2el_ms": [
3646.797
],
"mean_audio_rtf": [
0.1098
]
@@ -82,9 +79,6 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
"mean_e2el_ms": [
4059.0878
],
"mean_audio_rtf": [
0.1335
]
@@ -120,9 +114,6 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
"mean_e2el_ms": [
5639.5769
],
"mean_audio_rtf": [
0.1798
]
@@ -45,13 +45,6 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
"mean_e2el_ms": [
15793.1209,
30658.6632,
37463.4943,
47990.5839,
64980.2199
],
"mean_audio_rtf": [
0.169,
0.2041,
@@ -87,9 +80,6 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
"mean_e2el_ms": [
4556.4347
],
"mean_audio_rtf": [
0.138
]
@@ -123,9 +113,6 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
"mean_e2el_ms": [
4768.0972
],
"mean_audio_rtf": [
0.1609
]
@@ -161,9 +148,6 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
"mean_e2el_ms": [
6395.3517
],
"mean_audio_rtf": [
0.2048
]
+42 -1
View File
@@ -1,3 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for DFX runner metadata field exclusion."""
import json
@@ -105,7 +108,7 @@ def test_resolve_pytest_marks_hardware_dict_with_extra():
assert "H100" in names
assert "cuda" in names
assert "gpu" in names
assert "distributed_cuda" in names
assert "cards_2" in names
assert "full_model" in names
assert "diffusion" in names
assert "local_model" in names
@@ -512,3 +515,41 @@ def test_omni_duplex_expected_audio_turns_rejects_incomplete_session():
{"expected_duplex_audio_turns_per_session": 4},
1,
)
def test_omni_tpot_baseline_accepts_measured_finite_sample():
from tests.dfx.perf.scripts.run_benchmark import assert_result
assert_result(
{
"completed": 1,
"Hardware": "H100",
"num_tpot_samples": 1,
"mean_tpot_ms": 10.0,
},
{"baseline": {"H100": {"mean_tpot_ms": 20.0}}},
1,
)
@pytest.mark.parametrize(
("num_tpot_samples", "mean_tpot_ms", "match"),
[
(0, float("nan"), "no measurable TPOT samples"),
(1, float("nan"), "mean_tpot_ms is not finite"),
],
)
def test_omni_tpot_baseline_rejects_missing_or_nonfinite_sample(num_tpot_samples, mean_tpot_ms, match):
from tests.dfx.perf.scripts.run_benchmark import assert_result
with pytest.raises(AssertionError, match=match):
assert_result(
{
"completed": 1,
"Hardware": "H100",
"num_tpot_samples": num_tpot_samples,
"mean_tpot_ms": mean_tpot_ms,
},
{"baseline": {"H100": {"mean_tpot_ms": 20.0}}},
1,
)
@@ -101,7 +101,12 @@ class TestAttentionSpec:
def test_block_sparse_defaults_applied_when_backend_selected(self):
spec = AttentionSpec(backend="RAINFUSION_ATTN")
assert spec.block_sparse.sparsity == 0.8
assert spec.backend_kwargs() == {"sparsity": 0.8, "start_step": 0}
assert spec.backend_kwargs() == {
"sparsity": 0.8,
"start_step": 0,
"end_step": 0,
"precision": "bf16",
}
def test_block_sparse_skip_layers_selector_expanded(self):
spec = AttentionSpec(
@@ -112,6 +117,8 @@ class TestAttentionSpec:
assert spec.backend_kwargs() == {
"sparsity": 0.9,
"start_step": 12,
"end_step": 0,
"precision": "bf16",
"skip_layers": [0, 1, 2, 38],
}
@@ -121,12 +128,44 @@ class TestAttentionSpec:
@pytest.mark.parametrize(
"block_sparse",
[{"sparsity": 1.5}, {"start_step": -1}],
[{"sparsity": 1.5}, {"start_step": -1}, {"end_step": -1}, {"precision": "bogus"}],
)
def test_block_sparse_invalid_values(self, block_sparse):
with pytest.raises(ValueError):
AttentionSpec(backend="RAINFUSION_ATTN", block_sparse=block_sparse)
def test_block_sparse_end_step_precision_forwarded(self):
spec = AttentionSpec(
backend="RAINFUSION_ATTN",
block_sparse={"sparsity": 0.9, "start_step": 5, "end_step": 3, "precision": "fp8"},
)
assert spec.backend_kwargs() == {
"sparsity": 0.9,
"start_step": 5,
"end_step": 3,
"precision": "fp8",
}
def test_block_sparse_precision_enum_input_normalized(self):
"""RainFusionPrecision enum members must normalize to their string value.
Regression: the str,Enum members passed the value check but str() later
produced "RainFusionPrecision.MIX" instead of "mix" (StrEnum fixes it).
"""
from vllm_omni.diffusion.attention.backends.rainfusion_attn import RainFusionConfig
from vllm_omni.diffusion.data import RainFusionPrecision
spec = AttentionSpec(
backend="RAINFUSION_ATTN",
block_sparse={"sparsity": 0.8, "precision": RainFusionPrecision.MIX},
)
kwargs = spec.backend_kwargs()
assert kwargs["precision"] == "mix"
# The config path is where str(value) used to mangle enum members.
cfg = RainFusionConfig.from_backend_kwargs(kwargs)
assert cfg.precision == "mix"
assert type(cfg.precision) is str
class TestAttentionConfig:
def test_empty_config(self):
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Geometry contract RAINFUSION_ATTN enforces before handing a forward to rf_v2.
@@ -11,18 +11,25 @@ to the always-kept prefix segment before generating the block mask.
"""
import dataclasses
import sys
import types
from unittest import mock
import pytest
import torch
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
from vllm_omni.diffusion.attention.backends import rainfusion_attn
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout, VideoTokenSpan
from vllm_omni.diffusion.attention.backends.rainfusion_attn import (
_BLOCK_SIZE,
RainFusionAttentionBackend,
RainFusionAttentionImpl,
RainFusionPlan,
)
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
PREFIX_ROWS = 710 # 14 text rows + 696 audio rows
ALIGNED_GRID = (62, 24, 40) # 1280x768 -> 59520 video rows, 465 blocks
MISALIGNED_GRID = (62, 24, 42) # 1344x768 -> 62496 video rows, 488 blocks + 32 rows
@@ -97,6 +104,62 @@ def test_video_segment_must_be_the_tail_of_packed_document_zero():
assert make_impl()._resolve_plan(metadata) is None
def test_ref2va_multi_video_spans_resolve():
layout = VideoTokenLayout(
used_len=12000,
video_spans=(
VideoTokenSpan(start=128, latent_grid=(4, 16, 64), role="reference"),
VideoTokenSpan(start=5000, latent_grid=(4, 16, 64), role="target"),
),
)
plan = make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout))
assert plan is not None
assert plan.used_len == 12000
assert plan.video_spans == [
{"start": 128, "latent_shape": [4, 16, 64]},
{"start": 5000, "latent_shape": [4, 16, 64]},
]
def test_invalid_ref2va_video_spans_fall_back_to_dense():
layout = VideoTokenLayout(
used_len=12000,
video_spans=(
VideoTokenSpan(start=128, latent_grid=(4, 16, 64), role="reference"),
VideoTokenSpan(start=4000, latent_grid=(4, 16, 64), role="target"),
),
)
assert make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout)) is None
def test_validate_available_rejects_legacy_mindiesd(monkeypatch):
mindiesd = types.ModuleType("mindiesd")
def sparse_attention(query, key, value, **kwargs):
return query
mindiesd.sparse_attention = sparse_attention
monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
with pytest.raises(ValueError, match="video_spans"):
RainFusionAttentionBackend.validate_available()
def test_validate_available_accepts_new_mindiesd(monkeypatch):
mindiesd = types.ModuleType("mindiesd")
def sparse_attention(query, key, value, *, video_spans=None, **kwargs):
return query
mindiesd.sparse_attention = sparse_attention
monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
RainFusionAttentionBackend.validate_available()
@pytest.mark.parametrize("grid", [(4, 24, 40), (1, 24, 40)])
def test_short_video_stays_dense(grid):
assert make_impl()._resolve_plan(make_metadata(grid)) is None
@@ -169,7 +232,9 @@ def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd()
heads, head_dim = 4, 128
q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3))
impl = make_impl()
# Explicit bf16: the dense comparison exercises the plain sparse kernel,
# and older MindIE-SD releases do not accept precision= (gate would raise).
impl = make_impl(precision="bf16")
impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0)
plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid))
out = impl._forward_sparse_npu(q, k, v, plan)
@@ -182,3 +247,90 @@ def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd()
).transpose(1, 2)
error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
# end_step tail fallback: the last ``end_step`` denoise steps must stay dense.
def test_end_step_tail_window_stays_dense():
"""step_idx inside the last end_step window must resolve no sparse plan."""
impl = make_impl(end_step=3)
fc = mock.Mock()
fc.denoise_step_idx = 47
fc.total_denoise_steps = 50
with (
mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True),
mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc),
):
assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is None
def test_end_step_before_tail_window_still_sparse():
"""step_idx before the tail window must still resolve a sparse plan."""
impl = make_impl(end_step=3)
fc = mock.Mock()
fc.denoise_step_idx = 46
fc.total_denoise_steps = 50
with (
mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True),
mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc),
):
assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is not None
def test_end_step_zero_never_triggers_tail_fallback():
"""end_step=0 must not fall back even on the final denoise step."""
impl = make_impl(end_step=0)
fc = mock.Mock()
fc.denoise_step_idx = 49
fc.total_denoise_steps = 50
with (
mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True),
mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc),
):
assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is not None
# precision capability gate: a mindiesd without sparse_attention(precision=)
# must fail fast instead of silently running the BF16 path.
def _fake_mindiesd_module():
"""Install a minimal stand-in so the import inside _forward_sparse_npu succeeds."""
import types
fake = types.ModuleType("mindiesd")
fake.sparse_attention = lambda *args, **kwargs: None
return fake
def test_precision_non_bf16_requires_mindiesd_support():
"""precision != bf16 against a mindiesd lacking the kwarg must raise RuntimeError."""
import sys
impl = make_impl(precision="mix")
sys.modules["mindiesd"] = _fake_mindiesd_module()
try:
with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=False):
with pytest.raises(RuntimeError, match="requires MindIE-SD"):
impl._forward_sparse_npu(None, None, None, None)
finally:
sys.modules.pop("mindiesd", None)
def test_precision_non_bf16_passes_gate_when_supported():
"""precision != bf16 with a capable mindiesd must not raise from the gate."""
import sys
impl = make_impl(precision="mix")
sys.modules["mindiesd"] = _fake_mindiesd_module()
try:
with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=True):
# q/k/v shapes: [B, S, N, D]; plan geometry must match S.
q = torch.randn(1, 8, 4, 128)
plan = RainFusionPlan(prefix_len=0, used_len=8, latent_shape=[2, 2, 2])
# The gate must pass; the fake mindiesd returns None so no crash.
out = impl._forward_sparse_npu(q, q, q, plan)
assert out is None
finally:
sys.modules.pop("mindiesd", None)
+241 -1
View File
@@ -1,5 +1,5 @@
# 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 TeaCache extractor functions.
@@ -18,6 +18,7 @@ Currently implemented:
- TestMiniMaxH3Extractor: MiniMaxH3DiTModel extractor
"""
import math
from abc import ABC, abstractmethod
from types import SimpleNamespace
from unittest.mock import MagicMock, Mock, patch
@@ -650,6 +651,29 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
torch.testing.assert_close(actual_video, expected_video)
torch.testing.assert_close(actual_audio, expected_audio)
def test_extractor_uses_prepared_rope_table(
self,
minimax_h3_module,
sample_inputs,
monkeypatch,
):
"""TeaCache must consume the denoise-loop RoPE table without rebuilding it."""
rope_table = minimax_h3_module.prepare_rope_table(
sample_inputs["img_position_ids"],
seq_len=sample_inputs["x"].shape[1],
)
def unexpected_rope_build(*args, **kwargs):
raise AssertionError("TeaCache extractor rebuilt the prepared RoPE table")
monkeypatch.setattr(minimax_h3_module, "prepare_rope_table", unexpected_rope_build)
monkeypatch.setattr(minimax_h3_module.rope, "forward", unexpected_rope_build)
context = extract_minimax_h3_context(
minimax_h3_module,
**{**sample_inputs, "rope_table": rope_table},
)
context.run_transformer_blocks()
def test_run_transformer_blocks_forwards_packed_layout(
self,
minimax_h3_module,
@@ -705,6 +729,222 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
assert video_logits.shape[0] == sample_inputs["img_pos_info"]["position_ids"].shape[0]
assert audio_logits.shape[0] == sample_inputs["audio_pos_info"]["position_ids"].shape[0]
@pytest.mark.sp
def test_teacache_strict_sp_keeps_cache_state_rank_local(
self,
minimax_h3_module,
sample_inputs,
monkeypatch,
):
"""Strict SP must not mix local TeaCache state with gathered block rows."""
from vllm_omni.diffusion.attention.ops import minimax_h3_modulation
from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook
from vllm_omni.diffusion.distributed import parallel_state
seq_len = sample_inputs["x"].shape[1]
local_len = seq_len // 2
hidden_size = minimax_h3_module.hidden_size
video_width = minimax_h3_module.arch.latents_dim * math.prod(minimax_h3_module.arch.patch_size)
audio_width = minimax_h3_module.arch.audio_latents_dim
captured_modulation_indices: list[torch.Tensor] = []
class _StrictSPPrepare(nn.Module):
def __init__(self):
super().__init__()
self.global_indices: list[torch.Tensor] = []
def forward(self, hidden, rope_table, combined_indices):
assert hidden.shape[0] == local_len
assert rope_table.shape[0] == local_len
self.global_indices.append(combined_indices.clone())
return hidden, rope_table, combined_indices.narrow(0, 0, local_len)
class _StrictSPGather(nn.Module):
def __init__(self):
super().__init__()
self.inputs: list[torch.Tensor] = []
self.outputs: list[torch.Tensor] = []
def forward(self, tensor):
self.inputs.append(tensor.clone())
if tensor.shape[-1] == hidden_size:
raise AssertionError("TeaCache must not gather hidden states in strict SP")
assert tensor.shape == (local_len, video_width + audio_width)
# Model the all-gather of compact logits from every SP rank.
# postprocess() indexes the resulting global layout with
# infer_out_pos/audio_pos, including rows owned by another rank.
gathered = tensor.repeat(seq_len // local_len, 1)
assert gathered.shape == (seq_len, video_width + audio_width)
self.outputs.append(gathered)
return gathered
class _StrictSPGroup:
def __init__(self):
super().__init__()
self.local_decisions: list[int] = []
self.ops = []
def all_reduce(self, decision, op):
self.local_decisions.append(int(decision.item()))
self.ops.append(op)
# Simulate another rank requiring the second step to execute
# the collective block path even though this rank would cache.
if len(self.local_decisions) == 2:
decision.fill_(1)
return decision
class _StrictSPBlock(nn.Module):
def __init__(self):
super().__init__()
self.block_indices: list[torch.Tensor] = []
def adaln_proj(self, t_emb):
del t_emb
zeros = torch.zeros(3, hidden_size, dtype=torch.float32)
return zeros, zeros
def norm1(self, hidden):
return hidden.clone()
def forward(self, hidden, *, combined_indices, rope_table, **kwargs):
del kwargs
assert hidden.shape[0] == local_len
assert rope_table.shape[0] == local_len
self.block_indices.append(combined_indices.clone())
return hidden + 1
class _StrictSPFinalLayer(nn.Module):
def forward(self, hidden, *, t_emb, inverse_indices):
del t_emb, inverse_indices
assert hidden.shape == (local_len, hidden_size)
return hidden[:, :video_width].float(), hidden[:, :audio_width].float()
local_prepare = _StrictSPPrepare()
gather = _StrictSPGather()
sp_group = _StrictSPGroup()
block = _StrictSPBlock()
minimax_h3_module.local_sp_prepare = local_prepare
minimax_h3_module.sp_gather = gather
minimax_h3_module.blocks = nn.ModuleList([block])
minimax_h3_module.final_layer = _StrictSPFinalLayer()
monkeypatch.setattr(minimax_h3_module, "_rope_local_span", lambda _seq_len: (0, local_len))
monkeypatch.setattr(
minimax_h3_module,
"_embed",
lambda **kwargs: (
torch.arange(local_len * hidden_size, dtype=torch.float32).reshape(local_len, hidden_size),
torch.zeros(3, hidden_size, dtype=torch.float32),
),
)
monkeypatch.setattr(
minimax_h3_module,
"prepare_rope_table",
lambda *args, **kwargs: torch.zeros(local_len, 6, dtype=torch.float32),
)
monkeypatch.setattr(parallel_state, "get_sp_group", lambda: sp_group)
original_indexed_scale_shift = minimax_h3_modulation.indexed_scale_shift_
def capture_indexed_scale_shift(hidden, shift, scale, indices):
captured_modulation_indices.append(indices.clone())
return original_indexed_scale_shift(hidden, shift, scale, indices)
monkeypatch.setattr(minimax_h3_modulation, "indexed_scale_shift_", capture_indexed_scale_shift)
local_img_pos = torch.tensor([0, 1])
local_audio_pos = torch.tensor([2, 3])
token_tags = torch.full((seq_len,), -1, dtype=torch.long)
token_tags[local_img_pos] = 0
token_tags[local_audio_pos] = 2
inverse_indices = torch.zeros(seq_len, dtype=torch.long)
inputs = {
**sample_inputs,
"inverse_indices": inverse_indices,
"token_tags": token_tags,
"update_mask": torch.ones(local_img_pos.numel(), dtype=torch.bool),
"prompt_embeds": torch.empty(0, 6),
"img_pos_info": {"position_ids": local_img_pos},
"audio_pos_info": {"position_ids": local_audio_pos},
"text_pos_info": {"position_ids": torch.empty(0, dtype=torch.long)},
"img_pos_for_infer_output_info": {"position_ids": local_img_pos},
"refiner_packed_seq_params": {
"cu_seqlens_q": torch.tensor([0, 0], dtype=torch.int32),
"max_seqlen_q": 0,
},
}
expected_local_indices = (inverse_indices * 3 + token_tags.clamp(min=0))[:local_len]
hook = TeaCacheHook(
TeaCacheConfig(
transformer_type="MiniMaxH3DiTModel",
rel_l1_thresh=0.3,
)
)
hook.initialize_hook(minimax_h3_module)
hook.new_forward(minimax_h3_module, **inputs)
hook.new_forward(minimax_h3_module, **inputs)
video_logits, audio_logits = hook.new_forward(minimax_h3_module, **inputs)
state = hook.state_manager.get_state()
assert state.previous_residual is not None
assert state.previous_residual.shape == (local_len, hidden_size)
assert len(captured_modulation_indices) == 3
assert all(torch.equal(indices, expected_local_indices) for indices in captured_modulation_indices)
assert sp_group.local_decisions == [1, 0, 0]
assert len(sp_group.ops) == 3
assert all(op == torch.distributed.ReduceOp.MAX for op in sp_group.ops)
assert len(local_prepare.global_indices) == 2
assert all(indices.shape == (seq_len,) for indices in local_prepare.global_indices)
assert len(block.block_indices) == 2
assert all(torch.equal(indices, expected_local_indices) for indices in block.block_indices)
assert len(gather.inputs) == 3
assert video_logits.shape == (local_img_pos.numel(), video_width)
assert audio_logits.shape == (local_audio_pos.numel(), audio_width)
@pytest.mark.cpu
def test_teacache_non_strict_sp_gathers_hidden_once(
self,
minimax_h3_module,
sample_inputs,
monkeypatch,
):
"""The non-strict SP cache path must not gather global hidden rows twice."""
from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook
class _Prepare(nn.Module):
def forward(self, hidden, rope_table, combined_indices):
return hidden, rope_table, combined_indices
class _Gather(nn.Module):
def __init__(self):
super().__init__()
self.inputs: list[torch.Tensor] = []
def forward(self, tensor):
self.inputs.append(tensor.clone())
return tensor
seq_len = sample_inputs["x"].shape[1]
gather = _Gather()
minimax_h3_module.sp_prepare = _Prepare()
minimax_h3_module.sp_gather = gather
monkeypatch.setattr(minimax_h3_module, "_rope_local_span", lambda _seq_len: (0, seq_len))
hook = TeaCacheHook(
TeaCacheConfig(
transformer_type="MiniMaxH3DiTModel",
rel_l1_thresh=0.3,
)
)
hook.initialize_hook(minimax_h3_module)
hook.new_forward(minimax_h3_module, **sample_inputs)
assert len(gather.inputs) == 1
assert gather.inputs[0].shape == (seq_len, minimax_h3_module.hidden_size)
def test_invalid_module_raises_error(self):
"""Test that invalid module without blocks raises ValueError."""
invalid_module = Mock()
@@ -9,6 +9,7 @@ import torch
from diffusers import DiffusionPipeline # pyright: ignore[reportPrivateImportUsage]
from PIL import Image
import vllm_omni.diffusion.data as diffusion_data
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
from vllm_omni.diffusion.data import (
@@ -629,6 +630,46 @@ class TestPipelineArgumentsHandling:
with pytest.raises(ValueError):
pipeline.forward(DiffusionRequestBatch(requests=[problematic_request]))
interpolation_request = _make_request(
sampling_params=OmniDiffusionSamplingParams(
enable_frame_interpolation=True,
),
)
with pytest.raises(ValueError, match="Frame interpolation is not supported with the Diffusers backend"):
pipeline.forward(DiffusionRequestBatch(requests=[interpolation_request]))
def test_diffusers_config_preserves_checkpoint_class_for_metadata(self, mocker):
class MockWanImageToVideoPipeline:
pass
od_config = _make_od_config(model_class_name=None)
mocker.patch(
"vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index",
return_value={"_class_name": "WanImageToVideoPipeline"},
)
mocker.patch.object(
diffusion_data,
"diffusers",
SimpleNamespace(WanImageToVideoPipeline=MockWanImageToVideoPipeline),
)
mocker.patch.object(OmniDiffusionConfig, "update_multimodal_support")
od_config.enrich_config()
assert od_config.model_class_name == "WanImageToVideoPipeline"
assert od_config.diffusers_pipeline_cls is MockWanImageToVideoPipeline
def test_diffusers_config_falls_back_to_adapter_without_pipeline_index(self, mocker):
od_config = _make_od_config(model_class_name=None)
mocker.patch(
"vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index",
return_value=None,
)
od_config.enrich_config()
assert od_config.model_class_name == "DiffusersAdapterPipeline"
@pytest.mark.advanced_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@@ -0,0 +1,117 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
from dataclasses import dataclass
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.diffusion.distributed import a2a_permute
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_jit_build_includes_cuda_headers_from_nvidia_wheels(tmp_path, monkeypatch) -> None:
cu13_include = tmp_path / "nvidia" / "cu13" / "include"
nccl_include = tmp_path / "nvidia" / "nccl" / "include"
nccl_lib = tmp_path / "nvidia" / "nccl" / "lib" / "libnccl.so.2"
cu13_include.mkdir(parents=True)
nccl_include.mkdir(parents=True)
nccl_lib.parent.mkdir(parents=True)
(cu13_include / "cusparse.h").touch()
(nccl_include / "nccl.h").touch()
nccl_lib.touch()
load_kwargs = {}
monkeypatch.setattr(a2a_permute.sysconfig, "get_paths", lambda: {"purelib": str(tmp_path)})
monkeypatch.setattr(
torch.utils.cpp_extension,
"load",
lambda **kwargs: load_kwargs.update(kwargs),
)
monkeypatch.setattr(a2a_permute.symm_mem, "set_backend", lambda _backend: None)
monkeypatch.setattr(a2a_permute, "_BUILT", False)
a2a_permute.ensure_a2a_permute_available()
assert set(load_kwargs["extra_include_paths"]) == {str(cu13_include), str(nccl_include)}
assert load_kwargs["extra_ldflags"] == [str(nccl_lib)]
@dataclass
class _FakeAllocation:
size: int
device: torch.device
class _FakeHandle:
rank = 0
def __init__(self, allocation: _FakeAllocation) -> None:
self.allocation = allocation
self.views: list[tuple[tuple[int, ...], torch.dtype]] = []
def get_buffer(self, rank, shape, dtype):
assert rank == self.rank
self.views.append((shape, dtype))
return self.allocation, shape, dtype
def test_workspace_reuses_peak_capacity_across_shapes(monkeypatch) -> None:
allocations: list[_FakeAllocation] = []
all_reduce_calls: list[object] = []
synchronize_calls: list[torch.device] = []
stream_id = [1]
a2a_permute._SYMM_WORKSPACES.clear()
monkeypatch.setattr(a2a_permute.torch.cuda, "is_current_stream_capturing", lambda: False)
monkeypatch.setattr(
a2a_permute.torch.cuda,
"current_stream",
lambda _device: SimpleNamespace(cuda_stream=stream_id[0]),
)
monkeypatch.setattr(a2a_permute, "_resolve_process_group", lambda _name: object())
monkeypatch.setattr(a2a_permute.torch, "ones", lambda *_args, **_kwargs: object())
monkeypatch.setattr(
a2a_permute.dist,
"all_reduce",
lambda _tensor, group: all_reduce_calls.append(group),
)
monkeypatch.setattr(
a2a_permute.torch.accelerator,
"synchronize",
lambda device: synchronize_calls.append(device),
)
def empty(size, *, dtype, device):
assert dtype == torch.uint8
allocation = _FakeAllocation(size=size, device=torch.device(device))
allocations.append(allocation)
return allocation
monkeypatch.setattr(a2a_permute.symm_mem, "empty", empty)
monkeypatch.setattr(a2a_permute.symm_mem, "rendezvous", lambda allocation, _group: _FakeHandle(allocation))
device = torch.device("cuda:0")
first = a2a_permute._get_symm_buffer((2, 3), torch.float16, device, "group")
smaller = a2a_permute._get_symm_buffer((1, 4), torch.float16, device, "group")
larger = a2a_permute._get_symm_buffer((4, 3), torch.float16, device, "group")
assert first[0].size == 12
assert smaller[0] is first[0]
assert larger[0].size == 24
assert [allocation.size for allocation in allocations] == [12, 24]
assert len(a2a_permute._SYMM_WORKSPACES) == 1
assert len(all_reduce_calls) == 1
assert synchronize_calls == [device, device]
stream_id[0] = 2
with pytest.raises(RuntimeError, match="single CUDA stream"):
a2a_permute._get_symm_buffer((1, 4), torch.float16, device, "group")
a2a_permute.clear_a2a_permute_workspaces()
assert not a2a_permute._SYMM_WORKSPACES
assert synchronize_calls == [device, device, device]
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
@@ -194,6 +194,47 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]
@pytest.mark.cpu
@pytest.mark.core_model
@pytest.mark.parametrize(
("fully_shard_degree", "error_message"),
[(0, "fully_shard_degree must be positive"), (3, "must be divisible by fully_shard_degree")],
)
def test_invalid_hsdp_shard_size_fails_before_group_creation(monkeypatch, fully_shard_degree: int, error_message: str):
"""Reject invalid HSDP shard sizes before allocating any process groups."""
world_size = 4
created_groups: list[object] = []
def fail_if_group_created(*args, **kwargs):
del args, kwargs
created_groups.append(object())
raise AssertionError("a process group was created before HSDP validation")
fake_world_group = SimpleNamespace(
rank_in_group=0,
local_rank=0,
device_group=object(),
)
monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
monkeypatch.setattr(parallel_state, "init_model_parallel_group", fail_if_group_created)
for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_HSDP_REPLICATE", "_EXPERT_PARALLEL_GROUP_RANKS"):
monkeypatch.setattr(parallel_state, name, None)
for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
with pytest.raises(ValueError, match=error_message):
parallel_state.initialize_model_parallel(
fully_shard_degree=fully_shard_degree,
use_hsdp=True,
backend="gloo",
)
assert created_groups == []
@pytest.mark.cpu
@pytest.mark.core_model
def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for the Sequence Parallelism (SP) framework.
These tests verify the SP plan mechanism and hooks work correctly without
@@ -17,6 +17,7 @@ import pytest
import torch
import torch.nn as nn
from tests.helpers.mark import hardware_marks
from vllm_omni.diffusion.distributed.sp_plan import (
SequenceParallelInput,
SequenceParallelOutput,
@@ -419,7 +420,6 @@ class TestHookRegistration:
# =============================================================================
@pytest.mark.L4
class TestModelSpPlans:
"""Test that model _sp_plan definitions are valid.
@@ -427,6 +427,8 @@ class TestModelSpPlans:
May require GPU for model imports.
"""
pytestmark = hardware_marks(res={"cuda": "L4"})
def test_zimage_transformer_sp_plan(self):
"""Test ZImageTransformer2DModel _sp_plan structure.
@@ -9,6 +9,7 @@ Uses ``Tongyi-MAI/Z-Image-Turbo`` as the parity model.
from __future__ import annotations
import gc
import os
import time
from pathlib import Path
@@ -75,8 +76,13 @@ def _run_generate(
if num_requests < 2:
raise ValueError("num_requests must be >= 2 (1 warmup + >=1 timed)")
gc.collect()
current_omni_platform.empty_cache()
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.start()
try:
@@ -113,7 +119,10 @@ def _run_generate(
pass
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
finally:
monitor.stop()
+32
View File
@@ -0,0 +1,32 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
import torch
import torch.nn.functional as F
from tests.helpers.mark import hardware_test
from vllm_omni.diffusion.layers.activation import SiluAndMul
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
@pytest.mark.cpu
def test_silu_and_mul_native_matches_packed_reference() -> None:
packed = torch.randn(257, 256, dtype=torch.float32)
gate, up = packed.chunk(2, dim=-1)
torch.testing.assert_close(SiluAndMul().forward_native(packed), F.silu(gate) * up)
@hardware_test(res={"npu": "A3"}, num_cards=1)
def test_silu_and_mul_npu_matches_packed_reference() -> None:
packed = torch.randn(257, 256, device="npu", dtype=torch.bfloat16)
gate, up = packed.chunk(2, dim=-1)
torch.testing.assert_close(
SiluAndMul()(packed),
F.silu(gate) * up,
atol=2e-2,
rtol=2e-2,
)
@@ -0,0 +1,194 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import sys
import types
import pytest
import torch
import torch.nn.functional as F
from tests.helpers.mark import hardware_test
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
_HEAD_DIM = 128
_ROTARY_DIM = 96
_EPS = 1e-5
def _reference(
q: torch.Tensor,
k: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
rope_table: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
half = _ROTARY_DIM // 2
def apply(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
x = F.rms_norm(x, (_HEAD_DIM,), weight, _EPS)
cos = rope_table[..., :half].unsqueeze(1)
sin = rope_table[..., half:].unsqueeze(1)
return torch.cat(
(
x[..., :half] * cos - x[..., half:_ROTARY_DIM] * sin,
x[..., half:_ROTARY_DIM] * cos + x[..., :half] * sin,
x[..., _ROTARY_DIM:],
),
dim=-1,
)
return apply(q, q_weight), apply(k, k_weight)
def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
torch.manual_seed(7)
tokens, heads = 11, 4
q = torch.randn(tokens, heads, _HEAD_DIM, dtype=torch.bfloat16)
k = torch.randn(tokens, heads, _HEAD_DIM, dtype=torch.bfloat16)
q_weight = torch.randn(_HEAD_DIM, dtype=torch.bfloat16)
k_weight = torch.randn(_HEAD_DIM, dtype=torch.bfloat16)
freqs = torch.randn(tokens, _ROTARY_DIM // 2)
rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1).to(torch.bfloat16)
return q, k, q_weight, k_weight, rope_table
@pytest.mark.cpu
def test_npu_qk_norm_rope_uses_torch_npu_fused_primitives_without_mindiesd(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from vllm_omni.diffusion.layers import fused_qk_norm_rope as fused
calls: dict[str, int] = {"rms_norm": 0, "rotary_mul": 0}
def npu_rms_norm(x: torch.Tensor, gamma: torch.Tensor, epsilon: float):
calls["rms_norm"] += 1
return (F.rms_norm(x, (_HEAD_DIM,), gamma, epsilon),)
def npu_rotary_mul(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
calls["rotary_mul"] += 1
assert x.shape[-1] == _ROTARY_DIM
assert cos.shape == (1, x.shape[1], 1, _ROTARY_DIM)
assert sin.shape == cos.shape
half = _ROTARY_DIM // 2
return torch.cat(
(
x[..., :half] * cos[..., :half] - x[..., half:] * sin[..., :half],
x[..., half:] * cos[..., half:] + x[..., :half] * sin[..., half:],
),
dim=-1,
)
monkeypatch.setitem(
sys.modules,
"torch_npu",
types.SimpleNamespace(npu_rms_norm=npu_rms_norm, npu_rotary_mul=npu_rotary_mul),
)
monkeypatch.setattr(fused, "find_spec", lambda _: None)
q, k, q_weight, k_weight, rope_table = _inputs()
actual_q, actual_k = fused._npu_qk_norm_rope(
q,
k,
q_weight,
k_weight,
rope_table,
_EPS,
_ROTARY_DIM,
)
expected_q, expected_k = _reference(q, k, q_weight, k_weight, rope_table)
assert calls == {"rms_norm": 2, "rotary_mul": 2}
torch.testing.assert_close(actual_q, expected_q, atol=0, rtol=0)
torch.testing.assert_close(actual_k, expected_k, atol=0, rtol=0)
@pytest.mark.cpu
def test_npu_qk_norm_rope_uses_mindiesd_and_normalizes_packed_layout(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from vllm_omni.diffusion.layers import fused_qk_norm_rope as fused
calls: dict[str, int] = {"rms_norm": 0, "mindiesd_rope": 0}
def npu_rms_norm(x: torch.Tensor, gamma: torch.Tensor, epsilon: float):
calls["rms_norm"] += 1
return (F.rms_norm(x, (_HEAD_DIM,), gamma, epsilon),)
def rotary_position_embedding(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, **kwargs):
calls["mindiesd_rope"] += 1
assert x.dim() == 4
assert x.shape[0] == 1
assert x.shape[-1] == _ROTARY_DIM
assert cos.shape == (1, x.shape[1], 1, _ROTARY_DIM)
assert sin.shape == cos.shape
assert kwargs == {
"rotated_mode": "rotated_half",
"head_first": False,
"fused": True,
}
half = _ROTARY_DIM // 2
return torch.cat(
(
x[..., :half] * cos[..., :half] - x[..., half:] * sin[..., :half],
x[..., half:] * cos[..., half:] + x[..., :half] * sin[..., half:],
),
dim=-1,
)
monkeypatch.setitem(sys.modules, "torch_npu", types.SimpleNamespace(npu_rms_norm=npu_rms_norm))
monkeypatch.setitem(
sys.modules,
"mindiesd",
types.SimpleNamespace(rotary_position_embedding=rotary_position_embedding),
)
monkeypatch.setattr(fused, "find_spec", lambda _: object())
q, k, q_weight, k_weight, rope_table = _inputs()
actual_q, actual_k = fused._npu_qk_norm_rope(
q,
k,
q_weight,
k_weight,
rope_table,
_EPS,
_ROTARY_DIM,
)
expected_q, expected_k = _reference(q, k, q_weight, k_weight, rope_table)
assert calls == {"rms_norm": 2, "mindiesd_rope": 2}
torch.testing.assert_close(actual_q, expected_q, atol=0, rtol=0)
torch.testing.assert_close(actual_k, expected_k, atol=0, rtol=0)
@hardware_test(res={"npu": "A3"}, num_cards=1)
def test_fused_qk_norm_rope_npu_matches_reference(monkeypatch: pytest.MonkeyPatch) -> None:
"""Exercise the public Ascend dispatch with real torch_npu/MindIE operators."""
from vllm_omni.diffusion.layers import fused_qk_norm_rope as fused
q, k, q_weight, k_weight, rope_table = (tensor.to("npu") for tensor in _inputs())
calls = 0
original_npu_impl = fused._npu_qk_norm_rope
def capture_npu_impl(*args, **kwargs):
nonlocal calls
calls += 1
return original_npu_impl(*args, **kwargs)
monkeypatch.setattr(fused, "_npu_qk_norm_rope", capture_npu_impl)
actual_q, actual_k = fused.fused_qk_norm_rope(
q,
k,
q_weight,
k_weight,
rope_table,
_EPS,
head_dim=_HEAD_DIM,
rotary_dim=_ROTARY_DIM,
)
expected_q, expected_k = _reference(q, k, q_weight, k_weight, rope_table)
assert calls == 1
torch.testing.assert_close(actual_q, expected_q, atol=2e-2, rtol=2e-2)
torch.testing.assert_close(actual_k, expected_k, atol=2e-2, rtol=2e-2)
@@ -273,12 +273,15 @@ def test_required_hwr_miss_fails_before_ordinary_loading_or_publication(
assert loader.take_host_weight_plan() is None
def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
def _make_dlo_online_quant_config(dp_size: int = 2) -> OmniDiffusionConfig:
return OmniDiffusionConfig(
model="",
dtype=torch.float32,
quantization_config="fp8",
parallel_config=DiffusionParallelConfig(data_parallel_size=2),
parallel_config=DiffusionParallelConfig(
data_parallel_size=dp_size,
sequence_parallel_size=1,
),
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
)
@@ -824,6 +827,93 @@ def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch):
assert loader.take_host_weight_plan() is None
def test_dlo_allgather_online_int8_uses_ordinary_loader(monkeypatch):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
od_config = _make_dlo_online_quant_config()
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
model.transformer.quant_method = object.__new__(NPUInt8OnlineLinearMethod)
model.transformer.quant_method.uses_meta_device = True
calls: list[object] = []
allowlist_models: list[nn.Module] = []
original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
allowlist_models.append(candidate)
return original_allowlist_check(candidate)
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
loader.load_weights = ( # type: ignore[method-assign]
lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
)
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(
None,
"online quantization requires the ordinary loader",
),
)
assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
assert allowlist_models == [model]
assert calls == [("load", True), "process"]
assert loader.take_host_weight_plan() is None
def test_dlo_online_quant_group_size_one_skips_allgather_gate(monkeypatch):
"""A DLO group of one runs no weight collective, so an otherwise
unvalidated online method must not be rejected by the AllGather gate."""
class UnsupportedOnlineMethod:
uses_meta_device = True
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
od_config = _make_dlo_online_quant_config(dp_size=1)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
model.transformer.quant_method = UnsupportedOnlineMethod()
calls: list[object] = []
allowlist_models: list[nn.Module] = []
original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
allowlist_models.append(candidate)
return original_allowlist_check(candidate)
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
loader.load_weights = ( # type: ignore[method-assign]
lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
)
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(
None,
"online quantization requires the ordinary loader",
),
)
assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
assert allowlist_models == []
assert calls == [("load", True), "process"]
def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
@@ -847,10 +937,54 @@ def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
),
)
with pytest.raises(ValueError, match="per-tensor FP8 linears"):
with pytest.raises(ValueError, match="per-tensor FP8, INT8, and MXFP8 linears"):
loader.load_model(load_device="cpu")
def test_dlo_allgather_online_mxfp8_uses_ordinary_loader(monkeypatch):
mxfp8_config = pytest.importorskip("vllm_omni.quantization.mxfp8_config")
NPUMxfp8OnlineLinearMethod = mxfp8_config.NPUMxfp8OnlineLinearMethod
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
od_config = _make_dlo_online_quant_config()
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
model.transformer.quant_method = object.__new__(NPUMxfp8OnlineLinearMethod)
model.transformer.quant_method.uses_meta_device = True
calls: list[object] = []
allowlist_models: list[nn.Module] = []
original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
allowlist_models.append(candidate)
return original_allowlist_check(candidate)
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
loader.load_weights = ( # type: ignore[method-assign]
lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
)
loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(
None,
"online quantization requires the ordinary loader",
),
)
assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
assert allowlist_models == [model]
assert calls == [("load", True), "process"]
assert loader.take_host_weight_plan() is None
def test_hsdp_processes_quantized_weights_before_sharding(mocker):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
from vllm_omni.diffusion.offloader.module_collector import PipelineModules
@@ -1,11 +1,12 @@
# 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."""
import json
from dataclasses import replace
from types import SimpleNamespace
from typing import Any
import pytest
import torch
@@ -170,7 +171,10 @@ def _resolve_request_inputs_for_test(
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)
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):
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(
LTX25_DISTILLED_COMPONENT_PROFILE,
text_encoder_cls=object,
@@ -196,12 +212,19 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
vocoder_fallback_cls=None,
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(
model="org/converted-ltx25",
revision=revision,
dtype=torch.bfloat16,
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"))
@@ -227,6 +250,10 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
)
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):
calls["transformer_config"] = (model, subfolder, local_files_only, revision)
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.AutoTokenizer, "from_pretrained", fake_tokenizer)
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,
"LTX2VideoVaeNeighborhoodNattenProcessor",
lambda: natten_processor,
)
monkeypatch.setattr(
ltx2_components,
"create_transformer_from_config",
@@ -261,6 +294,14 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
"latent_upsampler",
}
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"] == (
od_config.model,
profile.transformer_subfolder,
@@ -270,6 +311,85 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
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):
(tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"}))
(tmp_path / "transformer").mkdir()
@@ -411,6 +531,46 @@ def test_ltx25_four_public_pipeline_semantics_are_disjoint():
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(
("pipeline_kind", "expected"),
[
@@ -673,7 +833,7 @@ def test_ltx25_distilled_two_stage_executes_custom_phase_schedules():
max_sequence_length=16,
)
prompt_context = object()
phase_calls = []
phase_calls: list[Any] = []
def resolve_request_inputs(req, **kwargs):
return request_inputs
@@ -1,5 +1,5 @@
# 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."""
@@ -258,6 +258,39 @@ class TestLTXImageToVideoConditioning:
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):
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
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])
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():
pipeline = _make_pipeline(LTX2Pipeline)
@@ -1,9 +1,10 @@
# 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."""
from types import SimpleNamespace
from typing import Any
import pytest
import torch
@@ -148,10 +149,13 @@ def test_ltx_phase_adapter_selects_execution_from_quantization(
):
pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config)
manifest = object()
captured = {}
captured: dict[str, Any] = {}
def resolve_artifact(*_args, **kwargs):
assert not kwargs
assert kwargs == {
"model_revision": None,
"artifact_revision": None,
}
return "adapter"
def parse(transformer, path):
@@ -1,5 +1,5 @@
# 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."""
@@ -8,6 +8,7 @@ import os
import tempfile
from dataclasses import replace
from types import SimpleNamespace
from typing import Any
import pytest
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"
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"
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)
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):
@@ -187,7 +213,41 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
expected.write_bytes(b"sidecar")
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(
@@ -900,7 +960,7 @@ def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
max_sequence_length=16,
)
prompt_context = object()
phase_calls = []
phase_calls: list[Any] = []
def resolve_request_inputs(req, **kwargs):
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_temporal_compression_ratio = 8
pipeline.latent_upsampler = FakeUpsampler()
activated_slots = []
activated_slots: list[Any] = []
if pipeline_cls is LTX2TwoStagePipeline:
pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append)
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(
("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():
phase_switches = []
phase_switches: list[Any] = []
ordinary_pipeline = object.__new__(LTX2TwoStagePipeline)
ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append)
for phase in LTX2_TWO_STAGE_RECIPE.phases:
+439 -3
View File
@@ -1,9 +1,10 @@
# 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 typing import Any
import pytest
import torch
@@ -36,6 +37,441 @@ def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch):
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:
@pytest.mark.parametrize(
("distributed_vae_state", "expected_decode_calls"),
@@ -138,7 +574,7 @@ class TestLTX23VaeDistributedDecode:
},
output_dtype=torch.float32,
)
seen = {}
seen: dict[str, Any] = {}
def exec_tile(task):
return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1))
@@ -38,6 +38,7 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
pipeline.default_audio_shift = 3.0
pipeline.device = torch.device("cpu")
pipeline.od_config = SimpleNamespace()
pipeline.text_encoder = object()
cache_spec = object()
quality_plan = SimpleNamespace(cache_dit=cache_spec)
pipeline._quality_policy = Mock()
@@ -110,6 +111,40 @@ def test_pipeline_import_registry_and_component_discovery():
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):
path.mkdir(parents=True)
(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.device = torch.device("cpu")
pipeline.od_config = SimpleNamespace()
pipeline.text_encoder = Mock()
pipeline._base_schedule_by_partition = schedules
pipeline._quality_policy = Mock()
pipeline._quality_policy.resolve.return_value = SimpleNamespace(cache_dit=None)
@@ -998,7 +1034,7 @@ def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch):
config.hidden_size,
bias=False,
)
attn_call = {}
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)
@@ -1020,7 +1056,7 @@ def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch):
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 = {}
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)
@@ -1095,7 +1131,7 @@ def test_encoder_forward_forwards_video_inputs():
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,
)
@@ -1681,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
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):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
@@ -2120,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():
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,
)
@@ -2195,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):
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,
)
@@ -0,0 +1,62 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from collections.abc import Iterator
import pytest
import torch
import torch.nn as nn
from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLEncoder
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@pytest.fixture(autouse=True)
def restore_cudnn_sdp_state() -> Iterator[None]:
original = torch.backends.cuda.cudnn_sdp_enabled()
try:
yield
finally:
torch.backends.cuda.enable_cudnn_sdp(original)
def _encoder_without_weights() -> MiniMaxH3Qwen3VLEncoder:
encoder = MiniMaxH3Qwen3VLEncoder("unused", device=torch.device("cpu"), load_model=False)
encoder.text_model = nn.Linear(1, 1)
return encoder
@pytest.mark.parametrize("initial_state", [False, True])
def test_encode_ids_restores_cudnn_sdp_after_success(monkeypatch: pytest.MonkeyPatch, initial_state: bool) -> None:
encoder = _encoder_without_weights()
expected = torch.tensor([1.0])
def encode(*args, **kwargs):
assert torch.backends.cuda.cudnn_sdp_enabled()
return expected
monkeypatch.setattr(encoder, "_encode", encode)
torch.backends.cuda.enable_cudnn_sdp(initial_state)
actual = encoder.encode_ids(torch.tensor([1]))
torch.testing.assert_close(actual, expected)
assert torch.backends.cuda.cudnn_sdp_enabled() is initial_state
@pytest.mark.parametrize("initial_state", [False, True])
def test_encode_ids_restores_cudnn_sdp_after_failure(monkeypatch: pytest.MonkeyPatch, initial_state: bool) -> None:
encoder = _encoder_without_weights()
def fail(*args, **kwargs):
assert torch.backends.cuda.cudnn_sdp_enabled()
raise RuntimeError("encode failed")
monkeypatch.setattr(encoder, "_encode", fail)
torch.backends.cuda.enable_cudnn_sdp(initial_state)
with pytest.raises(RuntimeError, match="encode failed"):
encoder.encode_ids(torch.tensor([1]))
assert torch.backends.cuda.cudnn_sdp_enabled() is initial_state
@@ -223,6 +223,8 @@ def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = set()
pipeline._lora_sigma_schedules = {}
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
@@ -321,6 +323,8 @@ def test_only_an_active_recognized_turbo_adapter_restricts_ref2va(monkeypatch):
pipeline.partition = "combined"
pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = set()
pipeline._lora_sigma_schedules = {}
request = _request("generic-peft")
def load():
@@ -0,0 +1,860 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
from types import SimpleNamespace
import pytest
import torch
from safetensors.torch import save_file
from vllm.lora.lora_weights import PackedLoRALayerWeights
from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
from vllm_omni.diffusion.models.minimax_h3.npu import lora as lora_module
from vllm_omni.diffusion.models.minimax_h3.npu.lora import load_minimax_h3_native_lora
from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
from vllm_omni.errors import OmniClientError
from vllm_omni.lora.request import LoRARequest
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
_TINY_HIDDEN = 2
_TINY_FFN = 2
_TINY_TIME = 2
_TINY_BLOCK_ADALN = 4
_TINY_FINAL_ADALN = 3
_TINY_RANK = 4
_TINY_NUM_QUERY_GROUPS = 2
_TINY_HEADS_PER_GROUP = 1
_TINY_HEAD_DIM = 1
_TINY_QKV_SLICE = _TINY_NUM_QUERY_GROUPS * _TINY_HEADS_PER_GROUP * _TINY_HEAD_DIM
# Grouped qkv stores (heads_per_group + 2) * head_dim rows per query group, so
# the attention inner size follows from the head layout and cannot be picked
# independently: 3 * inner must equal the grouped row count.
_TINY_INNER = _TINY_QKV_SLICE
@pytest.fixture(autouse=True)
def _use_tiny_native_dimensions(monkeypatch):
monkeypatch.setattr(lora_module, "_NATIVE_RANK", _TINY_RANK)
monkeypatch.setattr(lora_module, "_NATIVE_ALPHA", _TINY_RANK)
monkeypatch.setattr(lora_module, "_NATIVE_HIDDEN_SIZE", _TINY_HIDDEN)
monkeypatch.setattr(lora_module, "_NATIVE_ATTENTION_INNER_SIZE", _TINY_INNER)
monkeypatch.setattr(lora_module, "_NATIVE_FFN_HIDDEN_SIZE", _TINY_FFN)
monkeypatch.setattr(lora_module, "_NATIVE_TIME_EMBED_DIM", _TINY_TIME)
monkeypatch.setattr(lora_module, "_NATIVE_BLOCK_ADALN_OUT", _TINY_BLOCK_ADALN)
monkeypatch.setattr(lora_module, "_NATIVE_FINAL_ADALN_OUT", _TINY_FINAL_ADALN)
monkeypatch.setattr(lora_module, "_NATIVE_NUM_QUERY_GROUPS", _TINY_NUM_QUERY_GROUPS)
monkeypatch.setattr(lora_module, "_NATIVE_HEADS_PER_GROUP", _TINY_HEADS_PER_GROUP)
monkeypatch.setattr(lora_module, "_NATIVE_HEAD_DIM", _TINY_HEAD_DIM)
monkeypatch.setattr(lora_module, "_NATIVE_QKV_SLICE", _TINY_QKV_SLICE)
monkeypatch.setattr(
lora_module,
"_NATIVE_TARGET_DIMS",
{
"attn.qkv_proj": (_TINY_HIDDEN, 3 * _TINY_INNER),
"attn.out_proj": (_TINY_INNER, _TINY_HIDDEN),
"mlp.fc1": (_TINY_HIDDEN, 2 * _TINY_FFN),
"mlp.fc2": (_TINY_FFN, _TINY_HIDDEN),
"adaln_proj.linear": (_TINY_TIME, _TINY_BLOCK_ADALN),
},
)
monkeypatch.setattr(lora_module, "_NATIVE_FINAL_ADALN_DIMS", (_TINY_TIME, _TINY_FINAL_ADALN))
monkeypatch.setattr(
lora_module,
"_NATIVE_EXPECTED_TARGETS",
frozenset(
[
*(
f"blocks.{block_index}.{suffix}"
for block_index in range(50)
for suffix in lora_module._NATIVE_TARGET_SUFFIXES
)
]
+ [
*(
f"token_refiner.blocks.{block_index}.{suffix}"
for block_index in range(2)
for suffix in lora_module._NATIVE_TOKEN_REFINER_SUFFIXES
)
]
+ ["final_layer.adaln_proj.linear"]
),
)
def _request(path) -> LoRARequest:
return LoRARequest(
lora_name="flashgen",
lora_int_id=7,
lora_path=str(path),
)
def _write_tiny_native(
path,
*,
omit_target: str | None = None,
shape_overrides: dict[str, tuple[int, int]] | None = None,
metadata: dict[str, str] | None = None,
) -> None:
rank = _TINY_RANK
tensors = {}
overrides = shape_overrides or {}
for target in sorted(lora_module._NATIVE_EXPECTED_TARGETS):
if target == omit_target:
continue
input_dim, output_dim = lora_module._native_target_dims(target)
raw_target = f"transformer.{target}"
a_name = f"{raw_target}.lora_A.default.weight"
b_name = f"{raw_target}.lora_B.default.weight"
tensors[a_name] = torch.ones(overrides.get(a_name, (rank, input_dim)))
if target.endswith(".attn.qkv_proj"):
grouped_rows = []
for group in range(_TINY_NUM_QUERY_GROUPS):
grouped_rows.extend(
[
torch.full((1, rank), float(group * 10 + 1)),
torch.full((1, rank), float(group * 10 + 2)),
torch.full((1, rank), float(group * 10 + 3)),
]
)
tensors[b_name] = torch.cat(grouped_rows, dim=0)
elif target.endswith(".mlp.fc1"):
tensors[b_name] = torch.cat(
(
torch.full((output_dim // 2, rank), 2.0),
torch.full((output_dim // 2, rank), 1.0),
),
dim=0,
)
else:
tensors[b_name] = torch.ones(overrides.get(b_name, (output_dim, rank)))
save_file(
tensors,
str(path),
metadata=metadata
or {
"format": "pt",
"key_format": "minimax-h3-native",
"qkv_layout": "grouped",
"lora_rank": str(rank),
"lora_alpha": str(rank),
"base_schedule": "1.0,0.7,0.4,0.15,0.0",
"tasks": "t2va",
},
)
def test_h3_native_loads_and_packs_qkv_and_fc1(tmp_path):
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
loaded = load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
assert loaded is not None
lora_model, peft_helper, schedule = loaded
assert peft_helper.r == _TINY_RANK
assert schedule.num_inference_steps == 4
assert len(lora_model.loras) == 259
assert "blocks.0.adaln_proj.linear" in lora_model.loras
assert "final_layer.adaln_proj.linear" in lora_model.loras
# The writer tags group g with q=g*10+1, k=g*10+2, v=g*10+3, so a correct
# grouped -> q/k/v permutation yields slices that are uniform per role.
qkv = lora_model.get_lora("blocks.0.attn.qkv_proj")
assert isinstance(qkv, PackedLoRALayerWeights)
torch.testing.assert_close(qkv.lora_b[0], torch.tensor([[1.0] * _TINY_RANK, [11.0] * _TINY_RANK]))
torch.testing.assert_close(qkv.lora_b[1], torch.tensor([[2.0] * _TINY_RANK, [12.0] * _TINY_RANK]))
torch.testing.assert_close(qkv.lora_b[2], torch.tensor([[3.0] * _TINY_RANK, [13.0] * _TINY_RANK]))
fc1 = lora_model.get_lora("blocks.0.mlp.fc1")
assert isinstance(fc1, PackedLoRALayerWeights)
torch.testing.assert_close(fc1.lora_b[0], torch.full((_TINY_FFN, _TINY_RANK), 2.0))
torch.testing.assert_close(fc1.lora_b[1], torch.full((_TINY_FFN, _TINY_RANK), 1.0))
def test_h3_native_rejects_bad_metadata_and_ref2va(tmp_path):
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path, metadata={"key_format": "minimax-h3-native", "qkv_layout": "runtime"})
with pytest.raises(ValueError, match="qkv_layout='grouped'"):
load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
# The exact-filename guard runs before the partition check, so the artifact
# has to keep the published name for this to reach the Ref2VA rejection.
valid_dir = tmp_path / "valid"
valid_dir.mkdir()
valid = valid_dir / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(valid)
with pytest.raises(ValueError, match="supports T2VA only"):
load_minimax_h3_native_lora(
partition="ref2va",
lora_request=_request(valid),
lora_path=valid,
dtype=torch.float32,
)
renamed = tmp_path / "renamed.safetensors"
_write_tiny_native(renamed)
with pytest.raises(ValueError, match="supports only"):
load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(renamed),
lora_path=renamed,
dtype=torch.float32,
)
def test_h3_native_rejects_truncated_artifact(tmp_path):
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path, omit_target="blocks.49.mlp.fc2")
with pytest.raises(ValueError, match="target set does not match"):
load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
def test_h3_native_rejects_schedule_with_wrong_interval_count(tmp_path):
"""A mislabeled schedule must fail closed instead of silently changing steps."""
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(
path,
metadata={
"format": "pt",
"key_format": "minimax-h3-native",
"qkv_layout": "grouped",
"lora_rank": str(_TINY_RANK),
"lora_alpha": str(_TINY_RANK),
"base_schedule": "1.0,0.5,0.0",
"tasks": "t2va",
},
)
with pytest.raises(ValueError, match="4-interval base_schedule"):
load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
@pytest.mark.parametrize(
("raw_schedule", "message"),
[
(None, "requires safetensors metadata base_schedule"),
("", "must not be empty"),
(",,", "must not be empty"),
("1.0,oops,0.15,0.0", "is malformed"),
("0.9,0.7,0.4,0.15,0.0", "is malformed"),
],
)
def test_h3_native_rejects_malformed_schedule_metadata(tmp_path, raw_schedule, message):
"""safetensors metadata is string-only, so the loader owns the parse contract."""
metadata = {
"format": "pt",
"key_format": "minimax-h3-native",
"qkv_layout": "grouped",
"lora_rank": str(_TINY_RANK),
"lora_alpha": str(_TINY_RANK),
"tasks": "t2va",
}
if raw_schedule is not None:
metadata["base_schedule"] = raw_schedule
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path, metadata=metadata)
with pytest.raises(ValueError, match=message):
load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
def test_h3_native_declared_file_fails_closed_on_invalid_metadata(tmp_path):
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path, metadata={"key_format": "other"})
with pytest.raises(ValueError, match="requires safetensors metadata"):
load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
def test_pipeline_native_schedule_and_task_validation(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = {7}
pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])}
pipeline._base_schedule_by_partition = {"fl2va": None}
sampling = SimpleNamespace(
lora_request=LoRARequest("flashgen", 7, "/tmp/native.safetensors"),
lora_scale=1.0,
num_inference_steps=4,
extra_args={},
)
assert pipeline._sigma_schedule_for_request(sampling, "t2va").num_inference_steps == 4
with pytest.raises(OmniClientError, match="num_inference_steps must be 4"):
pipeline._validate_native_sampling(
SimpleNamespace(lora_request=sampling.lora_request, num_inference_steps=5),
task="t2va",
)
with pytest.raises(OmniClientError, match="supports T2VA requests only"):
pipeline._resolve_task(
"fl2va",
{},
has_turbo_lora=False,
has_native_lora=True,
)
pipeline._base_schedule_by_partition = {"fl2va": DMD2SigmaSchedule.from_positions([1.0, 0.5, 0.0])}
with pytest.raises(OmniClientError, match="already pins base_schedule"):
pipeline._sigma_schedule_for_request(sampling, "t2va")
def test_pipeline_native_step_count_follows_adapter_schedule():
"""The accepted step count is the adapter's interval count, never a literal."""
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline._native_lora_adapter_ids = {7}
pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.5, 0.0])}
request = LoRARequest("flashgen", 7, "/tmp/native.safetensors")
pipeline._validate_native_sampling(
SimpleNamespace(lora_request=request, num_inference_steps=2),
task="t2va",
)
with pytest.raises(OmniClientError, match="num_inference_steps must be 2 or omitted, not 3"):
pipeline._validate_native_sampling(
SimpleNamespace(lora_request=request, num_inference_steps=3),
task="t2va",
)
with pytest.raises(OmniClientError, match="requires num_inference_steps=2"):
pipeline._validate_native_sampling(
SimpleNamespace(lora_request=request, num_inference_steps=4),
task="t2va",
)
@pytest.mark.parametrize("step_execution", [False, True])
def test_pipeline_native_omitted_step_count_follows_execution_mode(step_execution):
"""Only request mode can resolve an omitted count from the adapter schedule.
``StepScheduler`` reads ``num_inference_steps`` off the request at admission,
before this pipeline sees it, so under step execution an omitted value can
only disagree with the denoise loop. The advertised contract has to match.
"""
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline._native_lora_adapter_ids = {7}
pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])}
pipeline.od_config = SimpleNamespace(step_execution=step_execution)
sampling = SimpleNamespace(
lora_request=LoRARequest("flashgen", 7, "/tmp/native.safetensors"),
num_inference_steps=None,
)
if not step_execution:
# Request mode defaults to the adapter's own interval count.
pipeline._validate_native_sampling(sampling, task="t2va")
return
with pytest.raises(OmniClientError, match="explicit num_inference_steps=4 under step execution"):
pipeline._validate_native_sampling(sampling, task="t2va")
# The step-mode message must not offer omission it cannot accept.
with pytest.raises(OmniClientError, match=r"must be 4, not 5$"):
pipeline._validate_native_sampling(
SimpleNamespace(lora_request=sampling.lora_request, num_inference_steps=5),
task="t2va",
)
def test_pipeline_replaces_native_classification_after_reload(monkeypatch, tmp_path):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
request = _request(path)
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
enable_distributed_layerwise_offload=False,
)
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = set()
pipeline._lora_sigma_schedules = {}
loaded = pipeline._load_diffusion_lora_adapter(
lora_request=request,
lora_path=path,
dtype=torch.float32,
)
assert loaded is not None
assert request.lora_int_id in pipeline._native_lora_adapter_ids
assert request.lora_int_id in pipeline._lora_sigma_schedules
monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None)
monkeypatch.setattr(pipeline_module, "load_minimax_h3_native_lora", lambda **_: None)
assert pipeline._load_diffusion_lora_adapter(lora_request=request, lora_path=path, dtype=torch.float32) is None
assert request.lora_int_id not in pipeline._native_lora_adapter_ids
assert request.lora_int_id not in pipeline._lora_sigma_schedules
@pytest.mark.parametrize(
"offload_mode",
[
"model-level CPU offload (--enable-cpu-offload)",
"layerwise offload (--enable-layerwise-offload)",
],
)
def test_h3_native_rejects_offload_modes(tmp_path, offload_mode):
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
with pytest.raises(ValueError, match="does not support"):
load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
unsupported_offload_mode=offload_mode,
)
def test_h3_native_allows_distributed_layerwise_offload(monkeypatch):
"""DLO keeps LoRA A/B buffers resident, so native must not fail closed."""
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.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
enable_distributed_layerwise_offload=True,
)
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = set()
pipeline._lora_sigma_schedules = {}
captured: dict[str, object] = {}
schedule = DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])
def load_native(**kwargs):
captured.update(kwargs)
return object(), object(), schedule
monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None)
monkeypatch.setattr(pipeline_module, "load_minimax_h3_native_lora", load_native)
request = _request("flashgen")
loaded = pipeline._load_diffusion_lora_adapter(
lora_request=request,
lora_path="flashgen",
dtype=torch.bfloat16,
)
assert loaded is not None
assert captured["unsupported_offload_mode"] is None
assert request.lora_int_id in pipeline._native_lora_adapter_ids
assert pipeline._lora_sigma_schedules[request.lora_int_id] is schedule
def test_h3_native_qkv_reorder_matches_base_loader_contract():
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import _reorder_grouped_qkv_to_qkv
num_query_groups = 56
heads_per_group = 1
head_dim = 128
grouped = torch.arange(num_query_groups * (heads_per_group + 2) * head_dim, dtype=torch.float32)
grouped = grouped.reshape(num_query_groups, (heads_per_group + 2) * head_dim)
reordered = _reorder_grouped_qkv_to_qkv(
grouped.reshape(-1, 1),
num_query_groups=num_query_groups,
heads_per_group=heads_per_group,
head_dim=head_dim,
).reshape(-1)
q_size = num_query_groups * heads_per_group * head_dim
k_size = num_query_groups * head_dim
q, k, v = torch.split(reordered, [q_size, k_size, k_size])
assert q.reshape(num_query_groups, head_dim)[0, 0] == grouped[0, 0]
assert k.reshape(num_query_groups, head_dim)[0, 0] == grouped[0, head_dim]
assert v.reshape(num_query_groups, head_dim)[0, 0] == grouped[0, 2 * head_dim]
def test_native_target_pattern_matches_component_qualified_names():
"""The manager matches "<component>.<module>", so the pattern needs the prefix."""
import regex as re
pattern = lora_module._NATIVE_TARGET_PATTERN
matched = [
target
for target in lora_module._NATIVE_EXPECTED_TARGETS
if re.search(pattern, f"transformer.{target}") is not None
]
assert len(matched) == 259
assert re.search(pattern, "transformer.blocks.0.attn.q_proj") is None
assert re.search(pattern, "transformer.token_refiner.blocks.0.adaln_proj.linear") is None
def test_replace_layers_binds_every_native_target(monkeypatch):
"""Regression for bound=0/259: layer replacement must reach all 259 modules."""
from vllm.lora.layers import BaseLayerWithLoRA
from vllm.lora.peft_helper import PEFTHelper
from vllm_omni.diffusion.lora import manager as manager_module
class _StubLoRALayer(BaseLayerWithLoRA):
def __init__(self, base_layer):
super().__init__()
self.base_layer = base_layer
def _linear() -> torch.nn.Module:
return torch.nn.Linear(_TINY_HIDDEN, _TINY_HIDDEN)
def _block(*, with_adaln: bool) -> torch.nn.Module:
block = torch.nn.Module()
block.attn = torch.nn.Module()
block.attn.qkv_proj = _linear()
block.attn.out_proj = _linear()
block.mlp = torch.nn.Module()
block.mlp.fc1 = _linear()
block.mlp.fc2 = _linear()
if with_adaln:
block.adaln_proj = torch.nn.Module()
block.adaln_proj.linear = _linear()
return block
transformer = torch.nn.Module()
transformer.blocks = torch.nn.ModuleList([_block(with_adaln=True) for _ in range(50)])
transformer.token_refiner = torch.nn.Module()
transformer.token_refiner.blocks = torch.nn.ModuleList([_block(with_adaln=False) for _ in range(2)])
transformer.final_layer = torch.nn.Module()
transformer.final_layer.adaln_proj = torch.nn.Module()
transformer.final_layer.adaln_proj.linear = _linear()
pipeline = torch.nn.Module()
pipeline.transformer = transformer
monkeypatch.setattr(manager_module, "from_layer_diffusion", lambda *, layer, **_: _StubLoRALayer(layer))
manager = object.__new__(DiffusionLoRAManager)
manager.pipeline = pipeline
manager.dtype = torch.float32
manager.max_cached_adapters = 1
manager._lora_modules = {}
manager._max_lora_rank = 64
manager._resident_lora_device = None
peft_helper = PEFTHelper.from_dict(
{
"r": _TINY_RANK,
"lora_alpha": _TINY_RANK,
"target_modules": lora_module._NATIVE_TARGET_PATTERN,
}
)
manager._replace_layers_with_lora(peft_helper)
assert set(manager._lora_modules) == {f"transformer.{target}" for target in lora_module._NATIVE_EXPECTED_TARGETS}
def test_legacy_manager_uses_native_loader(tmp_path):
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
class _Pipeline:
def _load_diffusion_lora_adapter(self, **kwargs):
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
MiniMaxH3Pipeline,
)
pipeline = object.__new__(MiniMaxH3Pipeline)
pipeline.partition = "fl2va"
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
enable_distributed_layerwise_offload=False,
)
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = set()
pipeline._lora_sigma_schedules = {}
return pipeline._load_diffusion_lora_adapter(**kwargs)
manager = object.__new__(DiffusionLoRAManager)
manager.pipeline = _Pipeline()
manager.dtype = torch.float32
manager._expected_lora_modules = {"qkv_proj", "fc1", "adaln_proj.linear"}
lora_model, peft_helper = manager._load_adapter(_request(path))
assert lora_model.id == 7
assert peft_helper.lora_alpha == _TINY_RANK
assert len(lora_model.loras) == 259
def test_pipeline_schedule_inactive_when_scale_zero():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = {7}
pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])}
pipeline._base_schedule_by_partition = {"fl2va": None}
sampling = SimpleNamespace(
lora_request=LoRARequest("flashgen", 7, "/tmp/native.safetensors"),
lora_scale=0.0,
num_inference_steps=4,
)
assert pipeline._sigma_schedule_for_request(sampling, "t2va") is None
assert not pipeline._has_active_native_lora(sampling)
def test_pipeline_schedule_falls_back_after_eviction(monkeypatch, tmp_path):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
request = _request(path)
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=False,
enable_layerwise_offload=False,
enable_distributed_layerwise_offload=False,
)
pipeline._turbo_lora_adapter_ids = set()
pipeline._native_lora_adapter_ids = set()
pipeline._lora_sigma_schedules = {}
pipeline._base_schedule_by_partition = {"fl2va": None}
loaded = pipeline._load_diffusion_lora_adapter(
lora_request=request,
lora_path=path,
dtype=torch.float32,
)
assert loaded is not None
assert request.lora_int_id in pipeline._lora_sigma_schedules
monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None)
monkeypatch.setattr(pipeline_module, "load_minimax_h3_native_lora", lambda **_: None)
assert pipeline._load_diffusion_lora_adapter(lora_request=request, lora_path=path, dtype=torch.float32) is None
assert request.lora_int_id not in pipeline._lora_sigma_schedules
sampling = SimpleNamespace(
lora_request=request,
lora_scale=1.0,
num_inference_steps=4,
)
assert pipeline._sigma_schedule_for_request(sampling, "t2va") is None
def test_native_packed_qkv_slices_are_tp2_divisible():
assert lora_module._NATIVE_QKV_SLICE % 2 == 0
assert lora_module._NATIVE_FFN_HIDDEN_SIZE % 2 == 0
def test_lora_manager_activates_native_packed_qkv(tmp_path):
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
loaded = load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
assert loaded is not None
lora_model, _, _ = loaded
packed = lora_model.get_lora("blocks.0.attn.qkv_proj")
assert isinstance(packed, PackedLoRALayerWeights)
assert len(packed.lora_b) == 3
assert all(b.shape[0] == _TINY_QKV_SLICE for b in packed.lora_b)
class _DummyLoRALayer:
def __init__(self):
self.n_slices = 3
self.output_slices = (_TINY_QKV_SLICE, _TINY_QKV_SLICE, _TINY_QKV_SLICE)
self.set_calls: list[tuple[list, list]] = []
self.reset_calls = 0
def reset_lora(self, index: int):
self.reset_calls += 1
def set_lora(self, index: int, lora_a, lora_b):
self.set_calls.append((lora_a, lora_b))
layer = _DummyLoRALayer()
manager = object.__new__(DiffusionLoRAManager)
manager.pipeline = torch.nn.Module()
manager._lora_modules = {"blocks.0.attn.qkv_proj": layer}
manager._registered_adapters = {
lora_model.id: lora_model,
}
manager._active_adapter_id = None
manager._adapter_scales = {}
manager._activate_adapter(lora_model.id, scale=0.5)
assert manager._active_adapter_id == lora_model.id
assert len(layer.set_calls) == 1
lora_a_list, lora_b_list = layer.set_calls[0]
assert len(lora_a_list) == 3
assert len(lora_b_list) == 3
torch.testing.assert_close(lora_b_list[0], packed.lora_b[0] * 0.5)
torch.testing.assert_close(lora_b_list[1], packed.lora_b[1] * 0.5)
torch.testing.assert_close(lora_b_list[2], packed.lora_b[2] * 0.5)
def test_lora_manager_tp2_splits_native_packed_qkv_per_rank():
"""Each Q/K/V slice must be divisible by TP2 so rank-local layers can shard rows."""
tp_size = 2
slice_rows = lora_module._NATIVE_QKV_SLICE
assert slice_rows % tp_size == 0
local_rows = slice_rows // tp_size
class _TpQkvLoRALayer:
def __init__(self, tp_rank: int):
self.tp_rank = tp_rank
self.n_slices = 3
self.output_slices = (local_rows, local_rows, local_rows)
self.set_calls: list[tuple[list, list]] = []
def reset_lora(self, index: int):
return
def set_lora(self, index: int, lora_a, lora_b):
self.set_calls.append((lora_a, lora_b))
for tp_rank in range(tp_size):
layer = _TpQkvLoRALayer(tp_rank)
start = tp_rank * local_rows
end = start + local_rows
full_packed = PackedLoRALayerWeights(
module_name="blocks.0.attn.qkv_proj",
rank=lora_module._NATIVE_RANK,
lora_alphas=[64, 64, 64],
lora_a=[torch.ones(lora_module._NATIVE_RANK, lora_module._NATIVE_HIDDEN_SIZE)] * 3,
lora_b=[torch.full((slice_rows, lora_module._NATIVE_RANK), float(i + 1)) for i in range(3)],
scaling=[1.0, 1.0, 1.0],
)
local_b = [b[start:end].contiguous() for b in full_packed.lora_b]
layer.set_lora(0, full_packed.lora_a, local_b)
assert len(layer.set_calls) == 1
_, lora_b_list = layer.set_calls[0]
assert all(b.shape[0] == local_rows for b in lora_b_list)
assert lora_b_list[0][0, 0].item() == 1.0
assert lora_b_list[1][0, 0].item() == 2.0
assert lora_b_list[2][0, 0].item() == 3.0
# 259 modules, of which 52 qkv expand to 3 slices and 52 fc1 expand to 2.
_EXPECTED_NATIVE_SLICE_PAIRS = 415
def _iter_native_slice_pairs(lora_model):
"""Yield ``(module_name, lora_a, lora_b)`` per slice, flattening packed weights."""
for module_name, weights in lora_model.loras.items():
lora_a = weights.lora_a if isinstance(weights.lora_a, list) else [weights.lora_a]
lora_b = weights.lora_b if isinstance(weights.lora_b, list) else [weights.lora_b]
for a, b in zip(lora_a, lora_b, strict=True):
yield module_name, a, b
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
def test_native_lora_load_is_platform_agnostic(tmp_path, dtype):
"""The artifact is trained on Ascend, but loading it is not Ascend-only.
This whole module runs in the CPU job with no ``torch_npu`` present, so a
passing run is itself the evidence. The assertions pin the property that
makes it portable: tensors land on CPU in exactly the requested dtype and
the loader never picks a device, leaving placement to the pipeline.
"""
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
loaded = load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=dtype,
)
assert loaded is not None
lora_model, _, _ = loaded
tensors = {
f"{module_name}.{side}": tensor
for module_name, a, b in _iter_native_slice_pairs(lora_model)
for side, tensor in (("lora_a", a), ("lora_b", b))
}
assert sum(1 for _ in _iter_native_slice_pairs(lora_model)) == _EXPECTED_NATIVE_SLICE_PAIRS
assert {name: tensor.device.type for name, tensor in tensors.items() if tensor.device.type != "cpu"} == {}
assert {name: tensor.dtype for name, tensor in tensors.items() if tensor.dtype != dtype} == {}
@pytest.mark.cuda
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device")
def test_native_lora_delta_matches_between_cpu_and_cuda(tmp_path):
"""The packed adapter produces the same delta weight on CUDA as on CPU.
Covers the GPU side of the same artifact: every ``lora_b @ lora_a`` product,
including the reordered q/k/v slices, is recomputed on CUDA and compared
against the CPU result, so a GPU deployment cannot silently diverge.
"""
path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
_write_tiny_native(path)
loaded = load_minimax_h3_native_lora(
partition="fl2va",
lora_request=_request(path),
lora_path=path,
dtype=torch.float32,
)
assert loaded is not None
lora_model, _, _ = loaded
checked = 0
for module_name, a, b in _iter_native_slice_pairs(lora_model):
cpu_delta = b @ a
cuda_delta = b.cuda() @ a.cuda()
assert cuda_delta.device.type == "cuda"
torch.testing.assert_close(cuda_delta.cpu(), cpu_delta, msg=f"delta mismatch for {module_name}")
checked += 1
assert checked == _EXPECTED_NATIVE_SLICE_PAIRS
@@ -74,6 +74,8 @@ def test_t2va_and_fl2va_packing_keep_update_rows_separate():
assert fl2va["update_mask"].sum().item() == 12
assert (~fl2va["update_mask"]).sum().item() == 6
assert t2va["audio_pos"].numel() == fl2va["audio_pos"].numel() == 6
assert t2va["video_spans"] == ({"start": 10, "latent_grid": (2, 2, 3), "role": "target"},)
assert fl2va["video_spans"] == ({"start": 16, "latent_grid": (2, 2, 3), "role": "target"},)
def test_ref2va_packing_tracks_video_and_audio_update_masks():
@@ -101,6 +103,33 @@ def test_ref2va_packing_tracks_video_and_audio_update_masks():
assert packed["cu_seqlens"].tolist() == [0, 30, 64]
def test_ref2va_packing_publishes_physical_video_spans_only():
from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
minimax_h3_packed_sequence_ref2va_blocks,
)
packed = minimax_h3_packed_sequence_ref2va_blocks(
text_len=4,
latent_t=2,
latent_h=4,
latent_w=6,
audio_t=3,
ref_blocks=[
{"kind": "image", "latent_h": 4, "latent_w": 4},
{"kind": "video_audio", "ref_audio_t": 2, "latent_t": 2, "latent_h": 4, "latent_w": 6},
{"kind": "audio", "ref_audio_t": 1},
{"kind": "video", "ref_audio_t": 0, "latent_t": 1, "latent_h": 4, "latent_w": 4},
],
)
# Image and all audio rows intentionally stay out of the sparse spans.
assert packed["video_spans"] == (
{"start": 12, "latent_grid": (2, 2, 3), "role": "reference"},
{"start": 26, "latent_grid": (1, 2, 2), "role": "reference"},
{"start": 36, "latent_grid": (2, 2, 3), "role": "target"},
)
def test_condition_noise_is_seeded_and_keeps_clean_anchor_at_timestep_one():
from vllm_omni.diffusion.models.minimax_h3.condition_noise import (
minimax_h3_audio_cond_noise_aug_rows,
@@ -1,5 +1,5 @@
# 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 torch
@@ -113,6 +113,146 @@ def test_h3_rope_table_materializes_local_rows_in_fused_kernel_layout():
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
def test_h3_prepared_rope_table_is_rank_local_and_validated(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
class Rope(nn.Module):
def forward(self, position_ids):
freqs_half = position_ids[0].float()
return torch.cat((freqs_half, freqs_half), dim=-1)
model = object.__new__(h3.MiniMaxH3DiTModel)
nn.Module.__init__(model)
model.arch = h3.MiniMaxH3DiTArchConfig(rope_inv_freq_len=1)
model.rope = Rope()
model.local_sp_prepare = nn.Identity()
monkeypatch.setattr(h3, "_sequence_parallel_local_span", lambda *args, **kwargs: (1, 2))
position_ids = torch.arange(12, dtype=torch.long).reshape(1, 4, 3)
actual = model.prepare_rope_table(position_ids, seq_len=4)
expected = h3._build_rope_table(model.rope(position_ids[:, 1:3]))
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
model._validate_prepared_rope_table(actual, local_len=2, device=torch.device("cpu"))
with pytest.raises(ValueError, match="local_seq_len"):
model._validate_prepared_rope_table(actual[:1], local_len=2, device=torch.device("cpu"))
def test_denoise_branch_prepares_rope_table_once_per_npu_run(monkeypatch: pytest.MonkeyPatch):
from vllm_omni.diffusion.models.minimax_h3 import denoise_loop
from vllm_omni.diffusion.models.minimax_h3.denoise_loop import (
MiniMaxH3DenoiseBranch,
minimax_h3_denoise_loop,
)
monkeypatch.setattr(denoise_loop.current_omni_platform, "is_npu", lambda: True)
packed = {
"seq_len": torch.tensor(3),
"img_pos": torch.tensor([1]),
"audio_pos": torch.tensor([2]),
"text_pos": torch.tensor([0]),
"update_mask": torch.tensor([True]),
"cu_seqlens": torch.tensor([0, 3, 3], dtype=torch.int32),
"img_position_ids": torch.zeros(3, 3, dtype=torch.long),
"latent_grid": torch.tensor([1, 1, 1]),
"video_row_start": torch.tensor(1),
}
branch = MiniMaxH3DenoiseBranch(
packed=packed,
text_embeddings=torch.zeros(1, 2),
token_tags=torch.zeros(3, dtype=torch.long),
device=torch.device("cpu"),
)
class Model:
def __init__(self):
self.rope_table = torch.zeros(3, 6, dtype=torch.bfloat16)
self.prepare_calls = 0
self.forward_calls = 0
def prepare_rope_table(self, img_position_ids, *, seq_len):
assert img_position_ids is branch.static_kwargs["img_position_ids"]
assert seq_len == 3
self.prepare_calls += 1
return self.rope_table
def __call__(self, **kwargs):
assert kwargs["rope_table"] is self.rope_table
self.forward_calls += 1
return torch.zeros(1, 96), torch.zeros(1, 32)
model = Model()
minimax_h3_denoise_loop(
model=model,
positive=branch,
initial_video_rows=torch.zeros(1, 96),
initial_audio_rows=torch.zeros(1, 32),
keyframe_cond_rows=None,
sigmas_video=[1.0, 0.5, 0.0],
sigmas_audio=[1.0, 0.5, 0.0],
device=torch.device("cpu"),
)
assert model.prepare_calls == 1
assert model.forward_calls == 2
def test_denoise_branch_keeps_rope_construction_on_the_reference_path_off_npu(
monkeypatch: pytest.MonkeyPatch,
):
from vllm_omni.diffusion.models.minimax_h3 import denoise_loop
from vllm_omni.diffusion.models.minimax_h3.denoise_loop import (
MiniMaxH3DenoiseBranch,
minimax_h3_denoise_loop,
)
monkeypatch.setattr(denoise_loop.current_omni_platform, "is_npu", lambda: False)
packed = {
"seq_len": torch.tensor(3),
"img_pos": torch.tensor([1]),
"audio_pos": torch.tensor([2]),
"text_pos": torch.tensor([0]),
"update_mask": torch.tensor([True]),
"cu_seqlens": torch.tensor([0, 3, 3], dtype=torch.int32),
"img_position_ids": torch.zeros(3, 3, dtype=torch.long),
"latent_grid": torch.tensor([1, 1, 1]),
"video_row_start": torch.tensor(1),
}
branch = MiniMaxH3DenoiseBranch(
packed=packed,
text_embeddings=torch.zeros(1, 2),
token_tags=torch.zeros(3, dtype=torch.long),
device=torch.device("cpu"),
)
class Model:
def __init__(self):
self.forward_calls = 0
def prepare_rope_table(self, *_args, **_kwargs):
raise AssertionError("the reference path must not pre-build a rope_table")
def __call__(self, **kwargs):
assert "rope_table" not in kwargs
self.forward_calls += 1
return torch.zeros(1, 96), torch.zeros(1, 32)
model = Model()
minimax_h3_denoise_loop(
model=model,
positive=branch,
initial_video_rows=torch.zeros(1, 96),
initial_audio_rows=torch.zeros(1, 32),
keyframe_cond_rows=None,
sigmas_video=[1.0, 0.5, 0.0],
sigmas_audio=[1.0, 0.5, 0.0],
device=torch.device("cpu"),
)
assert model.forward_calls == 2
def test_strict_sp_local_span_uses_rank_owned_rows(monkeypatch):
from vllm_omni.diffusion import forward_context
from vllm_omni.diffusion.distributed import parallel_state
@@ -254,18 +254,3 @@ def test_pipeline_resolves_transformer_component_quant_config():
assert transformer_config.ignored_layers == ignored_layers
assert _resolve_component_quant_config(component_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
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from pathlib import Path
@@ -187,6 +187,7 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
"model": model_ref,
"enforce_eager": True,
"tensor_parallel_size": 2,
# The fused BF16 baseline only fits on H100-80GB with encoder TP.
"text_encoder_tp_size": 2,
"vae_use_tiling": True,
}
@@ -1,13 +1,18 @@
# 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 torch
import torch.nn.functional as F
from tests.helpers.mark import hardware_marks
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")
@@ -0,0 +1,385 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from unittest.mock import Mock
import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.gpu, pytest.mark.diffusion]
def _selected_operators():
from vllm_omni.diffusion.models.minimax_h3.ops.vae.dispatch import (
resolve_h3_vae_operators,
)
if not current_omni_platform.is_available():
pytest.skip("No accelerator is available")
device = current_omni_platform.get_torch_device()
operators = resolve_h3_vae_operators(device)
if operators is None:
pytest.skip("No H3 VAE operator implementation is registered for this target")
return device, operators
def _operator_set(*, supports=lambda _device: True):
from vllm_omni.diffusion.models.minimax_h3.ops.vae.dispatch import H3VAEOperatorSet
return H3VAEOperatorSet(
supports=supports,
qk_norm_rope=lambda *_args: None,
scaled_residual=lambda *_args: None,
)
def _qk_reference(x, cos, sin):
normalized = F.rms_norm(x.float(), (64,), None, 1e-5).to(x.dtype)
rotary = normalized[..., :48]
first, second = rotary.chunk(2, dim=-1)
rotated = torch.cat((-second, first), dim=-1)
return torch.cat(
(rotary * cos + rotated * sin, normalized[..., 48:]),
dim=-1,
)
def _vit_norm_input(_module, hidden_states):
return hidden_states.float()
def _same_dtype_norm_input(_module, hidden_states):
return hidden_states
def _failing_norm_input(_module, _hidden_states):
raise RuntimeError("unsupported remote normalization semantics")
@pytest.mark.parametrize(("batch", "sequence"), [(1, 1), (1, 195), (2, 1797)])
def test_h3_vae_qk_norm_rope_is_bit_exact(batch, sequence):
device, operators = _selected_operators()
torch.manual_seed(17)
qkv = torch.randn(
batch,
sequence,
32,
192,
device=device,
dtype=torch.float16,
)
q, k, _ = qkv.chunk(3, dim=-1)
cos = torch.randn(batch, sequence, 1, 48, device=device, dtype=torch.float16)
sin = torch.randn_like(cos)
expected_q = _qk_reference(q, cos, sin)
expected_k = _qk_reference(k, cos, sin)
with torch.inference_mode():
actual = operators.qk_norm_rope(q, k, (cos, sin), 1e-5)
assert actual is not None
actual_q, actual_k = actual
assert torch.equal(actual_q, expected_q)
assert torch.equal(actual_k, expected_k)
def test_h3_vae_scaled_residual_is_bit_exact():
device, operators = _selected_operators()
torch.manual_seed(29)
residual = torch.randn(195, 2048, device=device, dtype=torch.float32)
branch = torch.randn(195, 2048, device=device, dtype=torch.float16)
scale = torch.randn(2048, device=device, dtype=torch.float32)
expected = residual + branch * scale
with torch.inference_mode():
actual = operators.scaled_residual(residual, branch, scale)
assert actual is not None
assert torch.equal(actual, expected)
def test_h3_vae_exact_ops_reject_unsupported_inputs():
from vllm_omni.diffusion.models.minimax_h3.ops.vae.qk_norm_rope import (
try_qk_norm_rope_exact,
)
from vllm_omni.diffusion.models.minimax_h3.ops.vae.scaled_residual import try_scaled_residual_exact
q = torch.randn(1, 2, 32, 64, dtype=torch.float32)
cos = torch.randn(1, 2, 1, 48, dtype=torch.float32)
residual = torch.randn(2, 2048)
branch = torch.randn(2, 2048, dtype=torch.float16)
scale = torch.randn(2048)
with torch.inference_mode():
assert try_qk_norm_rope_exact(q, q, (cos, cos), 1e-5) is None
assert try_scaled_residual_exact(residual, branch, scale) is None
def _make_decoder():
class Attention(nn.Module):
def __init__(self):
super().__init__()
self.to_qkv = nn.Linear(8, 24)
self.to_out = nn.Linear(8, 8)
self.norm_q = nn.RMSNorm(8, elementwise_affine=False)
self.norm_k = nn.RMSNorm(8, elementwise_affine=False)
self.spatial_parallel = False
self.dim_head = 8
def perform_attention(self, query, _key, _value, _pack_info):
return query
def forward(self, hidden_states, rotary_pos_emb=None, pack_info=None):
return hidden_states
class FeedForward(nn.Module):
def __init__(self):
super().__init__()
self.w1 = nn.Linear(8, 32)
self.w2 = nn.Linear(16, 8)
self.use_gated = True
self.act_fn = nn.SiLU()
self._compile_forward_enabled = False
self._compile_forward_fatal = False
self._compiled_forward = None
def forward(self, hidden_states):
hidden_states = self.w1(hidden_states)
gate, hidden_states = hidden_states.chunk(2, dim=-1)
return self.w2(self.act_fn(gate) * hidden_states)
class Block(nn.Module):
def __init__(self):
super().__init__()
self.attn = Attention()
self.ff = FeedForward()
self.norm1 = nn.RMSNorm(8)
self.norm2 = nn.RMSNorm(8)
self.scale1 = nn.Parameter(torch.zeros(8))
self.scale2 = nn.Parameter(torch.zeros(8))
self.use_scale = True
def forward(self, hidden_states, rotary_pos_emb=None, pack_info=None):
return hidden_states
decoder = nn.Module()
decoder.transformer_blocks = nn.ModuleList([Block(), Block()])
decoder.proj_out = nn.Linear(8, 8)
return decoder
def test_video_vae_installs_exact_optimizations(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
class FakeModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.decoder = torch.nn.Module()
class FakeRemote(torch.nn.Module):
def __init__(self):
super().__init__()
self.model = FakeModel()
monkeypatch.setattr(
vae_module,
"_load_component_config",
lambda _path: {
"latent_channels": 1,
"latents_mean": [0.0],
"latents_std": [1.0],
},
)
monkeypatch.setattr(
vae_module,
"_load_remote_component",
lambda _path, _config: FakeRemote(),
)
install = Mock(return_value=True)
monkeypatch.setattr(vae_module, "install_h3_vae_optimizations", install)
monkeypatch.setattr(vae_module, "PinnedModuleStager", Mock())
video_vae = vae_module.MiniMaxH3VideoVAE(
"unused",
device=torch.device("cuda"),
load_device=torch.device("cpu"),
)
install.assert_called_once_with(
video_vae.model.decoder,
device=torch.device("cuda"),
)
def test_h3_vae_install_precasts_only_block_linears(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
operators = _operator_set()
monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: operators)
decoder = _make_decoder()
assert vae_ops.install_h3_vae_optimizations(
decoder,
device=torch.device("meta"),
)
for block in decoder.transformer_blocks:
assert block.attn.to_qkv.weight.dtype == torch.float16
assert block.attn.to_out.weight.dtype == torch.float16
assert block.ff.w1.weight.dtype == torch.float16
assert block.ff.w2.weight.dtype == torch.float16
assert block.forward.__func__.__name__ == "_optimized_transformer_block"
assert block.attn.forward.__func__.__name__ == "_optimized_attention"
assert block.ff.forward.__func__.__name__ == "_optimized_feed_forward"
assert block.attn._omni_qk_norm_rope is operators.qk_norm_rope
assert block._omni_scaled_residual is operators.scaled_residual
assert decoder.proj_out.weight.dtype == torch.float32
# Repeated installation is idempotent.
assert vae_ops.install_h3_vae_optimizations(
decoder,
device=torch.device("meta"),
)
def test_h3_vae_install_accepts_remote_integer_parallel_flag(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
monkeypatch.setattr(
vae_ops,
"resolve_h3_vae_operators",
lambda _device: _operator_set(),
)
decoder = _make_decoder()
for block in decoder.transformer_blocks:
# The official checkpoint loads this JSON boolean-like field as int 0.
block.attn.spatial_parallel = 0
assert vae_ops.install_h3_vae_optimizations(
decoder,
device=torch.device("meta"),
)
assert getattr(decoder, "_omni_h3_vae_optimizations_installed", False)
def test_h3_vae_swiglu_uses_post_linear_fp16_output(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
device, operators = _selected_operators()
monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: operators)
decoder = _make_decoder().to(device)
feed_forward = decoder.transformer_blocks[0].ff
reference_forward = type(feed_forward).forward
assert vae_ops.install_h3_vae_optimizations(
decoder,
device=device,
)
hidden_states = torch.randn(4, 8, device=device, dtype=torch.float32)
with torch.inference_mode(), torch.autocast(device.type, dtype=torch.float16):
expected = reference_forward(feed_forward, hidden_states)
actual = feed_forward(hidden_states)
assert expected.dtype == torch.float16
assert torch.equal(actual, expected)
def test_h3_vae_install_leaves_unsupported_target_untouched(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: None)
decoder = _make_decoder()
assert not vae_ops.install_h3_vae_optimizations(
decoder,
device=torch.device("meta"),
)
assert not hasattr(decoder, "_omni_h3_vae_optimizations_installed")
for block in decoder.transformer_blocks:
assert block.attn.to_qkv.weight.dtype == torch.float32
assert block.ff.w1.weight.dtype == torch.float32
assert block.forward.__func__.__name__ == "forward"
assert block.attn.forward.__func__.__name__ == "forward"
@pytest.mark.parametrize("norm_input", [None, _same_dtype_norm_input, _failing_norm_input])
def test_h3_vae_install_requires_fp32_attention_norm_semantics(monkeypatch, norm_input):
from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
monkeypatch.setattr(
vae_ops,
"resolve_h3_vae_operators",
lambda _device: _operator_set(),
)
decoder = _make_decoder()
forward_globals = type(decoder.transformer_blocks[0].attn).forward.__globals__
monkeypatch.setitem(forward_globals, "_vit_norm_input", norm_input)
assert not vae_ops.install_h3_vae_optimizations(
decoder,
device=torch.device("meta"),
)
assert not hasattr(decoder, "_omni_h3_vae_optimizations_installed")
assert decoder.transformer_blocks[0].attn.to_qkv.weight.dtype == torch.float32
@pytest.mark.parametrize(
"break_contract",
[
lambda block: delattr(block.attn, "spatial_parallel"),
lambda block: setattr(block.attn, "spatial_parallel", 2),
lambda block: setattr(block.attn, "perform_attention", None),
lambda block: setattr(block, "scale1", nn.Parameter(torch.zeros(7))),
lambda block: setattr(block.ff, "_compile_forward_enabled", True),
],
)
def test_h3_vae_install_rejects_incompatible_remote_contract(monkeypatch, break_contract):
from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
monkeypatch.setattr(
vae_ops,
"resolve_h3_vae_operators",
lambda _device: _operator_set(),
)
decoder = _make_decoder()
break_contract(decoder.transformer_blocks[0])
assert not vae_ops.install_h3_vae_optimizations(
decoder,
device=torch.device("meta"),
)
assert decoder.transformer_blocks[0].attn.to_qkv.weight.dtype == torch.float32
def test_h3_vae_dispatch_is_extended_by_adding_an_operator_set(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3.ops.vae import dispatch
first = _operator_set(supports=lambda _device: False)
added = _operator_set(supports=lambda device: device.type == "meta")
monkeypatch.setattr(dispatch, "H3_VAE_OPERATOR_TABLE", (first, added))
assert dispatch.resolve_h3_vae_operators(torch.device("meta")) is added
assert dispatch.resolve_h3_vae_operators(torch.device("cpu")) is None
def test_h3_vae_dispatch_selects_supported_cuda_capabilities(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3.ops.vae import dispatch
platform = Mock()
platform.is_cuda.return_value = True
platform.is_available.return_value = True
monkeypatch.setattr(dispatch, "HAS_TRITON", True)
monkeypatch.setattr(dispatch, "current_omni_platform", platform)
for capability in (90, 100, 103):
platform.get_device_capability.return_value.to_int.return_value = capability
assert dispatch.resolve_h3_vae_operators(torch.device("cuda:0")) is not None
for capability in (89, 101, 110):
platform.get_device_capability.return_value.to_int.return_value = capability
assert dispatch.resolve_h3_vae_operators(torch.device("cuda:0")) is None
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import json
from types import SimpleNamespace
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import gc
from typing import Any
import numpy as np
@@ -15,7 +16,11 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
AUDIO_MODEL: dict[str, dict[str, int | None]] = {
"stabilityai/stable-audio-open-1.0": {"cuda": 1500, "rocm": 1500},
# The inference peak includes backend workspaces as well as model weights.
# On ROCm, large MIOpen workspaces and allocator fragmentation can mask the
# resident-weight reduction, so use a conservative floor that still catches
# a disabled/no-op layerwise offloader.
"stabilityai/stable-audio-open-1.0": {"cuda": 1500, "rocm": 512},
}
IMAGE_VIDEO_MODELS: dict[str, dict[str, int | None]] = {
@@ -63,8 +68,9 @@ def run_inference(
) -> tuple[float, Any]:
current_omni_platform.empty_cache()
device_index = current_omni_platform.current_device()
monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
monitor.start()
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)
if model_name in AUDIO_MODEL:
params = AUDIO_MODEL_PARAMS
@@ -78,22 +84,37 @@ def run_inference(
# cache_backend="cache_dit",
**params["runner_params"],
) as runner:
# Measure steady-state inference memory, not model construction. Enabling
# layerwise offload first loads the model and then replaces each block's
# device storage with CPU-backed weights. Monitoring that transition
# captures both the original model and temporary staging allocations,
# which is not representative of layerwise-offloaded inference.
monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
current_omni_platform.reset_peak_memory_stats()
monitor.start()
# Refer to tests/e2e/offline_inference/test_wan22.py
# Use minimal settings for testing
output = runner.omni.generate(
"A cat sitting on a table",
OmniDiffusionSamplingParams(
generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
guidance_scale=1.0,
num_inference_steps=num_inference_steps,
**params["sampler_params"],
),
)
try:
# Refer to tests/e2e/offline_inference/test_wan22.py
# Use minimal settings for testing
output = runner.omni.generate(
"A cat sitting on a table",
OmniDiffusionSamplingParams(
generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
guidance_scale=1.0,
num_inference_steps=num_inference_steps,
**params["sampler_params"],
),
)
finally:
monitor.stop()
peak = monitor.peak_used_mb
monitor.stop()
# DeviceMemoryMonitor reports absolute device usage. Subtract this run's
# starting usage so process-wide compiler/workspace caches retained from a
# previous run do not make the second measurement order-dependent.
peak = max(0.0, monitor.peak_used_mb - initial_used_mb)
gc.collect()
current_omni_platform.empty_cache()
return peak, output
@@ -148,8 +169,10 @@ def test_layerwise_offload_diffusion_model(model_name: str):
assert expected_saved_memory is not None
# Verify that layerwise offloading significantly reduces memory usage
# Passes only if the actual savings exceeds the expected savings
assert layerwise_offload_peak_memory + expected_saved_memory < no_offload_peak_memory, (
# Passes only if the actual savings meets the expected savings
actual_saved_memory = no_offload_peak_memory - layerwise_offload_peak_memory
assert layerwise_offload_peak_memory + expected_saved_memory <= no_offload_peak_memory, (
f"Layerwise offload peak memory {layerwise_offload_peak_memory} MB "
f"should be significantly less than no offload peak memory {no_offload_peak_memory} MB"
f"should be at least {expected_saved_memory} MB less than no offload peak memory "
f"{no_offload_peak_memory} MB (actual savings: {actual_saved_memory} MB)"
)
@@ -245,6 +245,160 @@ class TestDistributedLayerwiseOffloadHook:
)
torch.testing.assert_close(rank0_block.weight_scale, expected_scale)
def test_allgather_reconstructs_online_int8_weight_and_scale(
self,
patched_offload_runtime,
monkeypatch,
):
"""DLO must reconstruct finalized online-INT8 tensors byte-for-byte.
One block carries both finalized online-INT8 layouts: the NPU method
leaves a contiguous pre-transposed (K, N) int8 weight with a 1-D fp32
per-channel scale, while the CUDA method leaves a transposed-view
weight (stride (1, K)) with a scalar fp32 scale. The int8 and fp32
dtype groups must survive the flat-shard AllGather independently.
"""
class OnlineInt8Block(nn.Module):
def __init__(self):
super().__init__()
logical = torch.arange(1.0, 13.0).reshape(3, 4)
# NPU online int8: qweight.t().contiguous() -> (K, N), 1-D scale
self.npu_weight = nn.Parameter(logical.t().contiguous().to(torch.int8), requires_grad=False)
self.npu_weight_scale = nn.Parameter(torch.tensor([2.0, 3.0, 4.0]), requires_grad=False)
# CUDA online int8: Parameter(qweight.t().data) -> transposed
# view (stride (1, K)) over the (N, K) storage, scalar scale
self.cuda_weight = nn.Parameter((logical + 100).to(torch.int8).t(), requires_grad=False)
self.cuda_weight_scale = nn.Parameter(torch.tensor([0.5]), requires_grad=False)
current_block = OnlineInt8Block()
rank0_block = OnlineInt8Block()
rank1_block = OnlineInt8Block()
expected = {name: tensor.detach().clone() for name, tensor in rank0_block.named_parameters()}
expected_strides = {name: tensor.stride() for name, tensor in rank0_block.named_parameters()}
def make_hook(block: OnlineInt8Block, rank: int) -> DistributedLayerwiseOffloadHook:
return DistributedLayerwiseOffloadHook(
next_block=block,
device=torch.device("cpu"),
dp_group=object(),
dp_size=2,
rank=rank,
copy_stream=DummyStream(),
comm_stream=DummyStream(),
pin_memory=False,
)
rank0_hook = make_hook(rank0_block, 0)
rank1_hook = make_hook(rank1_block, 1)
rank0_hook.initialize_hook(current_block)
rank1_hook.initialize_hook(OnlineInt8Block())
rank0_hook.gpu_shard_buffers = [
{dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
]
def fake_allgather(output, local_shard, *, group):
del group
remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
shard_size = local_shard.numel()
output[:shard_size].copy_(local_shard)
output[shard_size : 2 * shard_size].copy_(remote_shard)
monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
rank0_hook.prefetch_layer(slot=0, non_blocking=False)
for name, tensor in rank0_block.named_parameters():
assert tensor.stride() == expected_strides[name]
torch.testing.assert_close(
tensor.float(),
expected[name].float(),
rtol=0,
atol=0,
)
def test_allgather_reconstructs_online_mxfp8_weight_and_scale(
self,
patched_offload_runtime,
monkeypatch,
):
"""DLO must reconstruct finalized online-MXFP8 tensors byte-for-byte.
One block carries both finalized online-MXFP8 layouts: the NPU method
leaves a contiguous pre-transposed (K, N) fp8 weight with a contiguous
(K_groups/2, N, 2) e8m0 block scale, while the vLLM-kernel method
(XPU) leaves an (N, K) fp8 weight whose e8m0 scale is a transposed
view over a contiguous (K/32, N) buffer. The fp8 and e8m0 dtype
groups must survive the flat-shard AllGather independently.
"""
class OnlineMxfp8Block(nn.Module):
def __init__(self):
super().__init__()
logical = torch.arange(1.0, 13.0).reshape(3, 4)
# Build scales from raw e8m0 bit patterns (1 byte/elem).
npu_scale_bits = torch.arange(1, 13, dtype=torch.uint8).reshape(2, 3, 2)
kernel_scale_bits = torch.arange(1, 7, dtype=torch.uint8).reshape(2, 3)
# NPU online mxfp8: (N, K).t().contiguous() -> (K, N) fp8
# weight, contiguous (K_groups/2, N, 2) e8m0 scale
self.npu_weight = nn.Parameter(logical.t().contiguous().to(torch.float8_e4m3fn), requires_grad=False)
self.npu_weight_scale = nn.Parameter(npu_scale_bits.view(torch.float8_e8m0fnu), requires_grad=False)
# vLLM-kernel online mxfp8 (XPU): (N, K) fp8 weight, e8m0
# scale kept as a transposed view (stride (1, K/32)) over the
# contiguous (K/32, N) storage
self.kernel_weight = nn.Parameter((logical + 100).to(torch.float8_e4m3fn), requires_grad=False)
self.kernel_weight_scale = nn.Parameter(
kernel_scale_bits.view(torch.float8_e8m0fnu).t(), requires_grad=False
)
current_block = OnlineMxfp8Block()
rank0_block = OnlineMxfp8Block()
rank1_block = OnlineMxfp8Block()
expected = {name: tensor.detach().clone() for name, tensor in rank0_block.named_parameters()}
expected_strides = {name: tensor.stride() for name, tensor in rank0_block.named_parameters()}
def make_hook(block: OnlineMxfp8Block, rank: int) -> DistributedLayerwiseOffloadHook:
return DistributedLayerwiseOffloadHook(
next_block=block,
device=torch.device("cpu"),
dp_group=object(),
dp_size=2,
rank=rank,
copy_stream=DummyStream(),
comm_stream=DummyStream(),
pin_memory=False,
)
rank0_hook = make_hook(rank0_block, 0)
rank1_hook = make_hook(rank1_block, 1)
rank0_hook.initialize_hook(current_block)
rank1_hook.initialize_hook(OnlineMxfp8Block())
rank0_hook.gpu_shard_buffers = [
{dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
]
def fake_allgather(output, local_shard, *, group):
del group
remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
shard_size = local_shard.numel()
output[:shard_size].copy_(local_shard)
output[shard_size : 2 * shard_size].copy_(remote_shard)
monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
rank0_hook.prefetch_layer(slot=0, non_blocking=False)
for name, tensor in rank0_block.named_parameters():
assert tensor.stride() == expected_strides[name]
# Compare raw bytes in logical order: float8/e8m0 CPU upcasts are
# version-sensitive, while contiguous().view(uint8) is an exact
# byte-level comparison for both contiguous and transposed-view
# parameters.
torch.testing.assert_close(
tensor.detach().contiguous().view(torch.uint8),
expected[name].contiguous().view(torch.uint8),
rtol=0,
atol=0,
)
def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
"""DTensor wrapper should be preserved through prefetch/offload cycle."""
current_block = TinyBlock(_make_values(1.0))
@@ -81,7 +81,7 @@ class QualityTestConfig:
num_frames: int = 5 # only for t2v
seed: int = 42
gpu: str = "H100" # minimum GPU requirement
negative_prompt: str = ""
negative_prompt: str | None = ""
guidance_scale: float | None = None
sigmas: list[float] | None = None
enable_cpu_offload: bool = False
@@ -212,24 +212,14 @@ QUALITY_CONFIGS = [
quantization="fp8",
task="t2v",
prompt="A serene lakeside sunrise with mist over the water",
max_lpips=0.10,
height=256,
width=256,
num_frames=25,
num_inference_steps=8,
# Preserve the final scheduler trajectory used when this FP8 quality
# gate was established. Explicit LTX sigmas bypass dynamic shifting.
sigmas=[
1.0,
0.92185378074646,
0.8327768445014954,
0.7303033471107483,
0.6111654043197632,
0.4709382653236389,
0.30347853899002075,
0.10000002384185791,
0.0,
],
max_lpips=0.20,
height=384,
width=512,
num_frames=73,
num_inference_steps=10,
# Do not override the recipe's negative conditioning, guidance, or
# scheduler trajectory: this gate follows the supported LTX default.
negative_prompt=None,
),
]
@@ -296,8 +286,11 @@ def _generate_image(omni, config: QualityTestConfig):
).manual_seed(config.seed)
torch.accelerator.reset_peak_memory_stats()
prompt = {"prompt": config.prompt}
if config.negative_prompt is not None:
prompt["negative_prompt"] = config.negative_prompt
outputs = omni.generate(
{"prompt": config.prompt, "negative_prompt": config.negative_prompt},
prompt,
OmniDiffusionSamplingParams(
height=config.height,
width=config.width,
@@ -328,8 +321,11 @@ def _generate_video(omni, config: QualityTestConfig):
).manual_seed(config.seed)
torch.accelerator.reset_peak_memory_stats()
prompt = {"prompt": config.prompt}
if config.negative_prompt is not None:
prompt["negative_prompt"] = config.negative_prompt
outputs = omni.generate(
{"prompt": config.prompt, "negative_prompt": config.negative_prompt},
prompt,
OmniDiffusionSamplingParams(
height=config.height,
width=config.width,
@@ -465,6 +461,7 @@ def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch
def test_generate_video_forwards_sigmas(monkeypatch):
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
monkeypatch.setattr(current_omni_platform, "device_type", "cpu", raising=False)
@@ -475,7 +472,12 @@ def test_generate_video_forwards_sigmas(monkeypatch):
class DummyOmni:
def generate(self, _prompt, sampling):
captured.sampling = sampling
return [SimpleNamespace(images=[np.zeros((1, 2, 2, 3), dtype=np.float32)])]
return [
OmniRequestOutput.from_diffusion(
request_id="req",
images=[np.zeros((1, 2, 2, 3), dtype=np.float32)],
)
]
config = QualityTestConfig(
id="ltx-sigmas",
@@ -491,6 +493,41 @@ def test_generate_video_forwards_sigmas(monkeypatch):
assert captured.sampling.sigmas == [1.0, 0.5]
@pytest.mark.core_model
@pytest.mark.cpu
def test_ltx_quality_gate_uses_official_eager_defaults(monkeypatch):
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
monkeypatch.setattr(current_omni_platform, "device_type", "cpu", raising=False)
monkeypatch.setattr(torch.accelerator, "reset_peak_memory_stats", lambda: None, raising=False)
monkeypatch.setattr(torch.accelerator, "max_memory_allocated", lambda: 0, raising=False)
captured = SimpleNamespace(prompt=None, sampling=None)
class DummyOmni:
def generate(self, prompt, sampling):
captured.prompt = prompt
captured.sampling = sampling
return [
OmniRequestOutput.from_diffusion(
request_id="req",
images=[np.zeros((1, 2, 2, 3), dtype=np.float32)],
)
]
config = next(config for config in QUALITY_CONFIGS if config.id == "fp8_ltx2")
_generate_video(DummyOmni(), config)
assert (config.width, config.height, config.num_frames, config.num_inference_steps) == (512, 384, 73, 10)
assert config.max_lpips == 0.20
assert config.sigmas is None
assert config.guidance_scale is None
assert config.negative_prompt is None
assert "negative_prompt" not in captured.prompt
assert captured.sampling.sigmas is None
assert captured.sampling.guidance_scale is None
_marks = hardware_marks(res={"cuda": "H100"})
_OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUALITY_OUTPUT_DIR" in os.environ else None
@@ -501,8 +538,6 @@ def _quality_param(c: QualityTestConfig):
marks.append(
pytest.mark.skip(reason="Qwen-Image FP8 quality gate temporarily disabled (see CI / issue tracker).")
)
if c.id == "fp8_ltx2":
marks.append(pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/6245"))
return pytest.param(c, id=c.id, marks=marks)
+54 -1
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import asyncio
import queue
@@ -197,6 +197,59 @@ class TestRequestBatchCapability:
assert diffusion_engine_module.supports_request_batch(od_config) is True
def test_engine_rejects_multi_seq_diffusers_using_adapter_capability(
self,
monkeypatch: pytest.MonkeyPatch,
mocker: MockerFixture,
) -> None:
od_config = SimpleNamespace(
model_class_name="QwenImagePipeline",
custom_pipeline_args=None,
diffusion_load_format="diffusers",
streaming_output=False,
max_num_seqs=2,
)
registry_load = mocker.Mock(
side_effect=lambda model_class_name: (
_SingleRequestPipeline if model_class_name == "DiffusersAdapterPipeline" else _BatchCapablePipeline
)
)
monkeypatch.setattr(
diffusion_engine_module.DiffusionModelRegistry,
"_try_load_model_cls",
registry_load,
)
engine = DiffusionEngine.__new__(DiffusionEngine)
with pytest.raises(ValueError, match="max_num_seqs=1"):
engine._resolve_execution_mode(od_config)
registry_load.assert_called_once_with("DiffusersAdapterPipeline")
def test_engine_prefers_batch_capable_custom_pipeline_over_diffusers_adapter(
self,
monkeypatch: pytest.MonkeyPatch,
mocker: MockerFixture,
) -> None:
od_config = SimpleNamespace(
model_class_name="QwenImagePipeline",
custom_pipeline_args={"pipeline_class": _BatchCapablePipeline},
diffusion_load_format="diffusers",
streaming_output=False,
max_num_seqs=2,
)
registry_load = mocker.Mock(return_value=_SingleRequestPipeline)
monkeypatch.setattr(
diffusion_engine_module.DiffusionModelRegistry,
"_try_load_model_cls",
registry_load,
)
engine = DiffusionEngine.__new__(DiffusionEngine)
assert engine._resolve_execution_mode(od_config) is DiffusionExecutionMode.REQUEST_BATCH
registry_load.assert_not_called()
def test_supports_request_batch_uses_custom_pipeline_class(self, monkeypatch: pytest.MonkeyPatch) -> None:
od_config = SimpleNamespace(
model_class_name="SinglePipeline",
Loaded 100 of 344 files, more files were not shown because too many files have changed in this diff. Show more