Compare commits

..
Author SHA1 Message Date
Hongsheng Liu 48bd050974 docs(review): enforce configuration provenance
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-09 22:39:57 +08:00
237 changed files with 3690 additions and 20757 deletions

No files matched your search

+1 -1
View File
@@ -8,7 +8,7 @@
# NPU naming: {queue}_{resource_class}
# a2b3_npu_1 — ascend-a2b3, resource_class npu-1 (A2/B3 CI image)
# a2b3_npu_4 — ascend-a2b3, resource_class npu-4
# a3_npu_2 — ascend-a3, resource_class npu-2
# a3_npu_2 — ascend-a3, resource_class npu-2 (A3 CI image tag …:$BUILDKITE_COMMIT-a3)
x-npu-k8s-plugins: &npu_k8s_plugins
- kubernetes:
+4 -4
View File
@@ -139,7 +139,7 @@ steps:
timeout_in_minutes: 50
commands:
- pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: h100_1
mirror_hardwares: h100_2
- label: "Omni · MiniCPM-o 4.5 Duplex Test"
source_file_dependencies:
@@ -148,10 +148,10 @@ steps:
- tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
- tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
- tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
- tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
- tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
- tests/helpers/minicpmo_4_5_duplex.py
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
- vllm_omni/experimental/fullduplex/
- vllm_omni/model_executor/models/minicpmo_4_5/
@@ -164,7 +164,7 @@ steps:
timeout_in_minutes: 50
commands:
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: h100_1
mirror_hardwares: h100_2
- label: "TTS · VoxCPM2 Test"
source_file_dependencies:
+4 -43
View File
@@ -22,7 +22,7 @@ steps:
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"
mirror_hardwares: h100_1
mirror_hardwares: h100_2
- label: ":full_moon: Omni · Doc Test with L4"
timeout_in_minutes: 90
@@ -50,19 +50,6 @@ steps:
exit $$EXIT
mirror_hardwares: h100_2
- label: ":full_moon: Omni · MiniCPM-o 4.5 · Accuracy Test"
timeout_in_minutes: 180
commands:
- export SEED_TTS_WER_EVAL=1
- export SEED_TTS_EVAL_DEVICE=cuda:1
- |
set +e
pytest -s -v tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py -m "full_model and cuda" --run-level full_model
EXIT=$$?
buildkite-agent artifact upload "tests/e2e/accuracy/minicpmo_4_5/results/*.json"
exit $$EXIT
mirror_hardwares: h100_2
- label: ":full_moon: Omni · Perf Test · No Async Chunk"
key: nightly-omni-performance-no-async-chunk
timeout_in_minutes: 300
@@ -115,32 +102,6 @@ steps:
exit $$EXIT
mirror_hardwares: h100_3
- label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test"
key: nightly-omni-performance-minicpmo-4-5
timeout_in_minutes: 180
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
mirror_hardwares: h100_1
- label: ":full_moon: Omni · MiniCPM-o 4.5 · Duplex Seed-TTS Perf Test"
key: nightly-omni-performance-minicpmo-4-5-duplex-seed-tts
timeout_in_minutes: 180
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
mirror_hardwares: h100_1
- label: ":full_moon: Omni · Multi-Replica Startup Test with 4x H100"
timeout_in_minutes: 45
commands:
@@ -443,7 +404,7 @@ steps:
- label: ":full_moon: Diffusion X2V · LingBot Video MoE Function Test"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/online_serving/test_lingbot_video_moe.py -k "video_generation" -m "full_model and cuda" --run-level "full_model"
- pytest -s -v tests/e2e/online_serving/test_lingbot_video_moe.py -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2V · Doc Test"
@@ -466,10 +427,10 @@ steps:
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · MiniMax H3 I2VA/Ref2VA Accuracy Test"
- label: ":full_moon: Diffusion X2V · MiniMax H3 T2VA Accuracy Test"
timeout_in_minutes: 180
commands:
- pytest -s -v tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/minimax_h3/test_minimax_h3_t2va_similarity.py -m full_model --run-level full_model
mirror_hardwares: h100_4
- label: ":full_moon: Diffusion X2V · Other Accuracy Test"
+5 -5
View File
@@ -172,7 +172,7 @@ steps:
timeout 20m bash -c "
pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: h100_1
mirror_hardwares: h100_2
- label: "Omni · MiniCPM-o 4.5 Online Test"
source_file_dependencies:
@@ -183,7 +183,7 @@ steps:
timeout 20m bash -c "
pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5.py -m 'core_model' --run-level 'core_model'
"
mirror_hardwares: h100_1
mirror_hardwares: h100_2
- label: "Omni · MiniCPM-o 4.5 Duplex Test"
source_file_dependencies:
@@ -192,10 +192,10 @@ steps:
- tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
- tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
- tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
- tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
- tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
- tests/helpers/minicpmo_4_5_duplex.py
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
- vllm_omni/experimental/fullduplex/
- vllm_omni/model_executor/models/minicpmo_4_5/
@@ -210,7 +210,7 @@ steps:
timeout 20m bash -c "
pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'core_model and cuda' --run-level 'core_model'
"
mirror_hardwares: h100_1
mirror_hardwares: h100_2
- label: "TTS · VoxCPM2 Test"
source_file_dependencies:
+1 -29
View File
@@ -39,7 +39,7 @@ steps:
EXIT=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
exit $$EXIT
- label: ":full_moon: Omni · MiniCPM-o 4.5 Test"
- label: ":robot: Omni · MiniCPM-o 4.5 Test"
mirror_hardwares: a3_npu_2
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
@@ -47,34 +47,6 @@ steps:
HF_TOKEN: "${HF_TOKEN}"
commands:
- pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5_expansion.py -m 'full_model' --run-level 'full_model'
- label: ":full_moon: Omni · MiniCPM-o 4.5 · Duplex Test"
timeout_in_minutes: 180
mirror_hardwares: a3_npu_2
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
HF_TOKEN: "${HF_TOKEN}"
commands:
- >-
pytest -s -v
tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
-m "npu and A3 and omni"
--run-level full_model
- label: ":full_moon: Omni · MiniCPM-o 4.5 · Accuracy Test"
timeout_in_minutes: 180
mirror_hardwares: a3_npu_2
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
HF_TOKEN: "${HF_TOKEN}"
commands:
- |
set +e
pytest -s -v tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py -m "full_model and npu and A3" --run-level full_model
EXIT=$$?
buildkite-agent artifact upload "tests/e2e/accuracy/minicpmo_4_5/results/*.json"
exit $$EXIT
- group: ":card_index_dividers: TTS Model Test"
key: nightly-tts-test-group
depends_on: upload-nightly-pipeline
+12 -10
View File
@@ -3,20 +3,22 @@ steps:
key: ready-omni-group
depends_on: upload-ready-pipeline
steps:
- label: "Omni · MiniCPM-o 4.5 Test"
mirror_hardwares: a3_npu_2
- label: ":robot: Omni · MiniCPM-o 4.5 Test"
agents:
queue: "ascend-a3"
resource_class: "npu-2"
image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a3"
plugins:
- kubernetes:
podSpecPatch:
imagePullSecrets:
- name: swr-secret
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
HF_TOKEN: "${HF_TOKEN}"
commands:
- pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py -m 'core_model and npu and A3 and omni' --run-level 'core_model'
- label: "Omni · MiniCPM-o 4.5 Duplex Test"
mirror_hardwares: a3_npu_2
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
commands:
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'core_model and npu and A3 and omni' --run-level 'core_model'
- pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py -m 'core_model' --run-level 'core_model'
- group: ":card_index_dividers: TTS Model Test"
key: ready-tts-group
+28 -38
View File
@@ -1,9 +1,9 @@
# Repository Skills for vLLM-Omni
# Claude Skills for vLLM-Omni
This directory contains repository-scale skills maintained for `vllm-omni`.
They capture repeatable workflows for common contributor and maintainer tasks
such as model integration, CI-aligned testing, performance optimization, and
pull request review.
This directory contains Claude Code skills maintained for the `vllm-omni`
repository. These skills capture repeatable workflows for common contributor
tasks such as model integration, pull request review, and release note
generation.
## Directory Structure
@@ -14,41 +14,31 @@ include:
- `references/`: focused reference material used by the skill
- `scripts/`: small helper scripts used by the skill
## Using the Skills
Coding agents that discover repository skills can invoke the relevant skill by
name. For other agents, point them to the skill's `SKILL.md` and ask them to
read it before changing code. Combine a domain skill with `vllm-omni-test` when
tests or CI coverage change, then use `precheck-pr` for the contributor's final
self-check.
## Available Skills
- [`add-diffusion-model`](add-diffusion-model/SKILL.md): guides integration of
a new diffusion model into `vllm-omni`
- [`add-tts-model`](add-tts-model/SKILL.md): covers integration of new TTS
models and related serving workflows
- [`diffusion-perf-opt`](diffusion-perf-opt/SKILL.md): guides diffusion model
performance optimization, including profiling traces, parallel strategies,
stage timing analysis, and benchmark-driven tuning
- [`precheck-pr`](precheck-pr/SKILL.md): self-checks a branch before creating a
PR by validating title format, dead code, accuracy and performance claims,
and merge readiness
- [`quantization`](quantization/SKILL.md): guides quantization method selection,
model integration, checkpoint loading, and quality/performance validation
for vLLM-Omni
- [`review-pr`](review-pr/SKILL.md): provides a frozen-snapshot, contract-aware
workflow for maintainers and reviewers; contributors should use
`precheck-pr` for self-review
- [`vllm-omni-npu-model-runner-upgrade`](vllm-omni-npu-upgrade/SKILL.md):
upgrades NPU model runners to align with the latest vllm-ascend
`NPUModelRunner`
- [`vllm-omni-test`](vllm-omni-test/SKILL.md): guides generation and execution
of CI-aligned tests (L1–L4), pytest marker selection (`core_model` /
`advanced_model` / `full_model`, `omni` / `tts` / `diffusion`), Buildkite
wiring (`test-ready.yml`, `test-merge.yml`, `test-nightly.yml`,
`test-weekly.yml`), and copy-paste local plus CI-like `pytest` commands; see
`references/test-routing.md` for level-to-command mapping
- `add-diffusion-model`: guides integration of a new diffusion model into
`vllm-omni`
- `diffusion-perf-opt`: guides diffusion model performance optimization,
including profiling traces, parallel strategies, stage timing analysis, and
benchmark-driven tuning
- `quantization`: guides quantization method selection, model integration,
checkpoint loading, and quality/performance validation for vLLM-Omni
- `add-omni-model`: covers addition of new omni-modality model support
- `add-tts-model`: covers integration of new TTS models and related serving
workflows
- `generate-release-note`: helps prepare release notes for repository changes
- `precheck-pr`: self-check a branch before creating a PR — validates PR title
format, catches dead code, verifies accuracy/perf claims, and confirms merge
readiness
- `vllm-omni-test`: guides generation and execution of CI-aligned tests (L1–L4),
pytest marker selection (`core_model` / `advanced_model` / `full_model`,
`omni` / `tts` / `diffusion`), Buildkite wiring (`test-ready.yml`,
`test-merge.yml`, `test-nightly.yml`, `test-weekly.yml`), and copy-paste
local plus CI-like `pytest` commands; see `references/test-routing.md` for
level-to-command mapping
- `review-pr`: provides a structured workflow for reviewing pull requests
- `vllm-omni-npu-model-runner-upgrade`: upgrades NPU model runners to align with the
latest vllm-ascend NPUModelRunner
## Maintenance Guidelines
+8 -1
View File
@@ -94,6 +94,7 @@ If docs and live code disagree, verify the code/tests and report the drift.
| --- | --- |
| [model-addition-checklist.md](references/checks/model-addition-checklist.md) | A model, architecture, loader, processor, registry, or stage config is added. |
| [perf-verification.md](references/checks/perf-verification.md) | The PR makes a latency, throughput, memory, or quality claim. |
| [config-provenance.md](references/checks/config-provenance.md) | Changed or newly load-bearing behavior reads config, CLI, environment, legacy, or default values. |
| [test-quality-evaluation.md](references/checks/test-quality-evaluation.md) | Tests change, are absent for risky code, or may not exercise production behavior. |
| [tests-docs-checklist.md](references/checks/tests-docs-checklist.md) | Coverage, CI markers, examples, user docs, or PR evidence need review. |
| [verification.md](references/checks/verification.md) | Hardware, a server, or a runnable affected path is available for active verification. |
@@ -152,13 +153,19 @@ contract they protect or use only the applicable evidence checks.
Apply every category in [general-checks.md](references/process/general-checks.md) before
lower-priority comments.
For each changed value or behavior, trace:
For each changed or newly load-bearing value or behavior, trace:
```text
public ingress -> validation/defaulting -> producer -> transformations
-> stage/worker/connector boundary -> final consumer -> terminal cleanup
```
Treat unchanged context as in scope when the diff makes it authoritative for a
new consumer. When config, CLI, environment, legacy, or default sources can
control that value, apply
[config-provenance.md](references/checks/config-provenance.md) and prove the
documented source at the final sink with distinct sentinel values.
Cover every applicable offline/online, streaming/non-streaming, sync/async,
feature-on/off, topology, and compatibility path. Search bounded callers and
sibling implementations rather than assuming the changed hunk is the only path.
@@ -0,0 +1,49 @@
# Configuration Provenance
Read this reference when changed or newly load-bearing behavior reads config,
CLI, environment, legacy, alias, or default values. Always use it for endpoint,
host, port, address, connector metadata, and `*_extra_config` consumers, even
when the schema or documentation is unchanged.
## Trace the source of truth
For each semantic value, record one compact ledger:
```text
documented input -> parsed field -> competing sources and precedence
-> runtime projection/transport -> final consumer
```
- Search the exact field, environment variable, default literal, and aliases in
code, tests, docs, examples, and deployment files.
- Treat unchanged producers and context lines as review scope when the diff
starts relying on them or makes their output authoritative.
- Resolve the intended precedence from active branch-local docs, enforced
code/tests, and compatibility policy. Report a documented input that is
silently ignored or reinterpreted at a live sink.
- Check every supported construction path and feature-off behavior without
inventing parity for unsupported paths.
## Require discriminating evidence
Give every competing source a distinct value so the test reveals which source
wins. Prefer non-default sentinels; for example:
```text
documented config port = 25301
environment/default port = 8998
expected final endpoint = http://10.0.0.8:25301
```
Assert the final URL, connection target, runtime argument, or other live sink,
not only an intermediate object. A test that changes the environment value to
the expected configured value proves formatting, not configuration propagation.
For a bug fix, show the discriminating test fails on the frozen base and passes
on the reviewed head when the environment permits execution.
## Finding bar
A supported non-default value reaching the wrong endpoint, device, topology, or
runtime behavior is a correctness or compatibility finding. A missing
discriminating test alone is a validation gap unless repository policy requires
that regression coverage for the change.
@@ -16,6 +16,9 @@ were reverted or broken. Check:
instead of the changed behavior being mocked away;
- fakes preserve relevant types, MRO, shapes, devices, async behavior, and
lifecycle transitions;
- configuration tests give competing config, CLI, environment, legacy, and
default sources distinct non-default values and assert the final consumer;
setting only the source read by the implementation does not prove precedence;
- seeds, ordering, timing, synchronization, external services, and numeric
tolerances are deterministic or explicitly controlled;
- normal, invalid, boundary, feature-off, failure/cancellation, and regression
@@ -15,6 +15,9 @@ and [design documents](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/
- Trace each changed value through public ingress, validation/defaulting,
producer, transformations, process or device boundaries, final consumer, and
terminal cleanup.
- Treat unchanged code as in scope when the diff makes its value newly
load-bearing for a changed consumer; validate that value's origin and contract
rather than reviewing only added lines.
- Check every applicable sync/async, offline/online, streaming/non-streaming,
feature-off, topology, and compatibility path without demanding unsupported
modes merely for symmetry.
@@ -123,7 +126,9 @@ distinct live caller, invariant, or compatibility need.
## Finding bar
Anchor each finding to a changed `path:line`; name the trigger or call path,
current behavior, impact, and smallest fix direction. Treat pending CI, missing
hardware, or unsupported measurements as validation gaps unless the repository
contract makes that evidence a merge requirement.
Anchor each finding to the exact `path:line` that causes the behavior; it may be
unchanged when the diff makes it newly load-bearing, in which case name the
changed call path that brings it into scope. Include the trigger, impact, and
smallest fix direction. Treat pending CI, missing hardware, or unsupported
measurements as validation gaps unless the repository contract makes that
evidence a merge requirement.
@@ -166,6 +166,10 @@ index patch, worktree patch, and untracked-content manifest. Any mismatch makes
the review stale; discard the affected evidence and restart. Prefer one
root-cause comment to several symptom comments.
When the root cause is an unchanged line made newly load-bearing by the diff,
cite it in the consolidated review body instead of moving the comment to an
unrelated changed line merely to make it inline.
Local presentation is the default. Only explicit authorization permits GitHub
posting. When authorized, publish one consolidated final review after
validation; do not post preliminary or incremental comments. Submit a review
@@ -50,6 +50,7 @@ compatibility changes.
| --- | --- | --- |
| New or expanded model, loader, processor, registry, or stage config | [model-addition-checklist.md](../checks/model-addition-checklist.md) | [`add-tts-model`](../../../add-tts-model/SKILL.md) or [`add-diffusion-model`](../../../add-diffusion-model/SKILL.md) |
| Latency, throughput, memory, scaling, precision, or quality claim | [perf-verification.md](../checks/perf-verification.md) | [`diffusion-perf-opt`](../../../diffusion-perf-opt/SKILL.md) for diffusion |
| Config/CLI/environment/default precedence; endpoint, host, port, address, connector metadata, or `*_extra_config` | [config-provenance.md](../checks/config-provenance.md) | None |
| Tests changed, absent for risky behavior, or test-only | [test-quality-evaluation.md](../checks/test-quality-evaluation.md) | [`vllm-omni-test`](../../../vllm-omni-test/SKILL.md) |
| CI, examples, docs, public behavior, or contributor evidence | [tests-docs-checklist.md](../checks/tests-docs-checklist.md) | None |
| Suitable hardware/server and runnable affected path | [verification.md](../checks/verification.md) | None |
+30 -112
View File
@@ -1,11 +1,6 @@
# CODEOWNERS is used to automatically assign PR reviewers
# Format: path pattern @username or @team
# More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
#
# The vllm_omni/** source-tree rules below are derived from the `owners` and
# `primary_code_paths` front matter in docs/design/module/*.md. When updating
# ownership, update the design doc first, then regenerate the matching rule here.
# Last-match-wins: more specific rules must come AFTER broader ones.
# Default owners (fallback reviewers for all files)
* @Gaohan123 @hsliuustc0106 @ywang96
@@ -26,10 +21,25 @@
# CI
/.buildkite/ @yenuo26 @congw729 @NickCao
# ComfyUI / external apps
/apps/ @fhfuih @alex-jw-brooks
# Engine and model_executor
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks @NickCao
/vllm_omni/model_executor/ @tzhouam @gcanlin
/vllm_omni/worker/ @tzhouam @gcanlin
# Benchmarks (out-of-tree)
# Diffusion
/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
/vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
# Diffusion CPU offloading
/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao @xuechendi
# Omni and TTS support
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin @NickCao
# Benchmarks
/benchmarks/accuracy/ @david6666666
/benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi
@@ -37,112 +47,20 @@
/benchmarks/glm_image/ @princepride @JaredforReal
/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin
# =====================================================================
# Module design ownership (docs/design/module/*.md)
# Each section header cites the source design doc.
# =====================================================================
# Autoregressive runtime — docs/design/module/ar_runtime.md
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
/vllm_omni/worker/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
# Engine orchestration & stage runtime — docs/design/module/engine_orchestration.md, stage_runtime.md
/vllm_omni/engine/ @tzhouam @fake0fan
# File-level overrides: input/output modality contracts own these engine files
# — docs/design/module/input_output_modality_contracts.md
/vllm_omni/engine/messages.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/engine/serialization.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/engine/mm_outputs.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/engine/output_modality.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/engine/output_processor.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
# Input/output modality contracts — docs/design/module/input_output_modality_contracts.md
/vllm_omni/inputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/outputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/request.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/data_entry_keys.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
# Entrypoints and serving boundaries — docs/design/module/entrypoints.md
/vllm_omni/entrypoints/ @alex-jw-brooks @linyueqian @NickCao
# File-level overrides within entrypoints/
/vllm_omni/entrypoints/openai/errors.py @alex-jw-brooks @NickCao
/vllm_omni/entrypoints/cli/benchmark/ @alex-jw-brooks @Bounty-hunter
# Error contracts — docs/design/module/error_contracts.md
/vllm_omni/errors.py @alex-jw-brooks @NickCao
# OmniConnector (inter-stage transport) — docs/design/module/omni_connector.md
/vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
# Model integration — docs/design/module/model_integration.md
/vllm_omni/model_executor/ @tzhouam @gcanlin
/vllm_omni/model_executor/models/ @tzhouam @gcanlin @Sy0307 @amy-why-3459 @linyueqian
/vllm_omni/model_extras/ @tzhouam @gcanlin @princepride
/vllm_omni/plugins/ @gcanlin @tjtanaa @xuechendi
# Config override within model_executor — docs/design/module/vllm_omni_config.md
/vllm_omni/model_executor/stage_configs/ @lishunyang12 @alex-jw-brooks
# Execution platforms (hardware backends) — docs/design/module/execution_platforms.md
/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
/vllm_omni/attention/ @gcanlin @tjtanaa @xuechendi @david6666666 @lishunyang12
# Configuration & deploy — docs/design/module/vllm_omni_config.md
/vllm_omni/config/ @lishunyang12 @alex-jw-brooks
/vllm_omni/deploy/ @lishunyang12 @alex-jw-brooks
# Quantization — docs/design/module/quantization.md
/vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
# Observability — docs/design/module/observability.md
/vllm_omni/metrics/ @lishunyang12 @vraiti
# Profiling — docs/design/module/profiling.md
# Usability and Observability
/vllm_omni/deploy/ @tzhouam @lishunyang12
/vllm_omni/metrics/ @lishunyang12
/vllm_omni/profiler/ @gcanlin
# Benchmarking (in-tree) — docs/design/module/benchmarking.md
/vllm_omni/benchmarks/ @alex-jw-brooks @Bounty-hunter
# ComfyUI
/apps/ @fhfuih @alex-jw-brooks
# LoRA (no dedicated design doc; ownership intentionally narrowed)
/vllm_omni/lora/ @SamitHuang @wtomin
# Large Scale Serving
/vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
# Transformers utils (related path only in design docs; preserves prior ownership)
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
# Quantization and configuration
/vllm_omni/config/ @lishunyang12 @alex-jw-brooks
/vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
# =====================================================================
# Diffusion module family — docs/design/module/diffusion/
# Umbrella rule first; sub-module rules after it win via last-match.
# =====================================================================
# Diffusion umbrella (top-level files: data.py, forward_context.py, etc.)
# — docs/design/module/diffusion/index.md
/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
# Cache management — docs/design/module/cache_management.md
/vllm_omni/diffusion/cache/ @Isotr0py @princepride @SamitHuang
# Diffusion quantization — docs/design/module/quantization.md
/vllm_omni/diffusion/quantization/ @david6666666 @Isotr0py @lishunyang12
# Diffusion profiling — docs/design/module/profiling.md
/vllm_omni/diffusion/profiler/ @gcanlin
# Diffusion runtime (scheduler, executor, worker)
# — docs/design/module/diffusion/diffusion_runtime.md, continuous_batching.md
/vllm_omni/diffusion/sched/ @Isotr0py @princepride @SamitHuang @fhfuih
/vllm_omni/diffusion/executor/ @Isotr0py @princepride @SamitHuang @fhfuih
/vllm_omni/diffusion/worker/ @Isotr0py @princepride @SamitHuang @fhfuih
# Diffusion model integration (pipelines, registration, checkpoint loading)
# — docs/design/module/diffusion/diffusion_model_integration.md
/vllm_omni/diffusion/models/ @fhfuih @Bounty-hunter @wtomin
/vllm_omni/diffusion/model_loader/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
/vllm_omni/diffusion/lora/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
# Diffusion parallelism (rank topology, TP/SP/PP/VAE patch)
# — docs/design/module/diffusion/parallelism.md
/vllm_omni/diffusion/distributed/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
/vllm_omni/diffusion/attention/parallel/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
# Diffusion offloader (CPU offloading, DLO)
# — docs/design/module/diffusion/offloader.md
/vllm_omni/diffusion/offloader/ @yuanheng-zhao @david6666666 @lishunyang12
# Hardware plugins and accelerators
/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
+1 -4
View File
@@ -66,10 +66,7 @@ repos:
entry: python tools/pre_commit/check_tts_adapter.py
language: python
pass_filenames: false
files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$
# The budget-is-now-slack reminder is printed on a PASSING run, and
# pre-commit hides output from passing hooks unless the hook is verbose.
verbose: true
files: ^vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)$
# Keep `suggestion` last
- id: suggestion
-1
View File
@@ -45,7 +45,6 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
"hyperpyyaml>=1.2.0" \
"onnx>=1.16.0" \
"librosa>=0.10.0" \
"datasets>=3.0.0" \
--index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
--trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
--default-timeout=60 && \
-1
View File
@@ -46,7 +46,6 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
"onnx>=1.16.0" \
"librosa>=0.10.0" \
"torchmetrics>=1.9.0" \
"datasets>=3.0.0" \
--index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
--trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
--default-timeout=60 && \
+39 -35
View File
@@ -16,21 +16,31 @@ nav:
- Examples:
- examples/README.md
- Offline Inference:
- GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
- HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
- Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
- Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
- Speech-To-Video: user_guide/examples/offline_inference/speech_to_video.md
- InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md
- MiMo-Audio: user_guide/examples/offline_inference/mimo_audio.md
- Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
- Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
- Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
- Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
- Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
- X-To-Text: user_guide/examples/offline_inference/x_to_text.md
- X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
- Online Serving:
- BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
- vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
- Diffusers Backend Adapter: user_guide/examples/online_serving/diffusers_pipeline_adapter.md
- GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
- Image-To-Image: user_guide/examples/online_serving/image_to_image.md
- Image-To-Video: user_guide/examples/online_serving/image_to_video.md
- Speech-To-Video: user_guide/examples/online_serving/speech_to_video.md
- Text-To-Image: user_guide/examples/online_serving/text_to_image.md
- MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
- Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
- Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
- Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
- Text-To-Image: user_guide/examples/online_serving/text_to_image.md
- Text-To-Video: user_guide/examples/online_serving/text_to_video.md
- General:
- usage/*
@@ -40,30 +50,16 @@ nav:
- Models:
- models/supported_models.md
- Features:
- Overview: features/README.md
- Runtime and Stage Execution:
- Execution Modes and Streaming: user_guide/diffusion/execution_modes.md
- Sleep Mode: features/sleep_mode.md
- Quantization:
- Overview: user_guide/quantization/overview.md
- Online Quantization: user_guide/quantization/online.md
- Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
- FP8 W8A8: user_guide/quantization/fp8.md
- Int8 W8A8: user_guide/quantization/int8.md
- BitsAndBytes W4: user_guide/quantization/bitsandbytes.md
- MXFP8 W8A8: user_guide/quantization/mxfp8.md
- MXFP4 W4A4: user_guide/quantization/mxfp4.md
- ModelOpt: user_guide/quantization/modelopt.md
- AutoRound: user_guide/quantization/autoround.md
- msModelSlim: user_guide/quantization/msmodelslim.md
- Diffusion Acceleration:
- Overview and Compatibility:
- Overview: user_guide/diffusion_features.md
- Feature Compatibility: user_guide/feature_compatibility.md
- CPU Offloading: user_guide/diffusion/cpu_offload.md
- Sleep Mode: features/sleep_mode.md
- Diffusion Features:
- Overview: user_guide/diffusion_features.md
- Feature Compatibility: user_guide/feature_compatibility.md
- Regional Compilation: user_guide/diffusion/regional_compilation.md
- Cache Acceleration:
- TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
- Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
- Attention Backends: user_guide/diffusion/attention_backends.md
- Frame Interpolation: user_guide/diffusion/frame_interpolation.md
- Parallelism:
- Overview: user_guide/diffusion/parallelism/overview.md
- CFG Parallel: user_guide/diffusion/parallelism/cfg_parallel.md
@@ -72,15 +68,23 @@ nav:
- Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
- Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
- VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md
- Attention Backends: user_guide/diffusion/attention_backends.md
- Regional Compilation: user_guide/diffusion/regional_compilation.md
- Frame Interpolation: user_guide/diffusion/frame_interpolation.md
- Startup and Loading: user_guide/diffusion/startup_and_loading.md
- CPU Offloading: user_guide/diffusion/cpu_offload.md
- LoRA: user_guide/diffusion/lora.md
- Experimental:
- Session State Manager: features/session_state_manager.md
- Integrations:
- Custom Diffusion Pipeline: features/custom_pipeline.md
- Custom Pipeline: features/custom_pipeline.md
- Execution Modes: user_guide/diffusion/execution_modes.md
- Quantization:
- Quantization: user_guide/quantization/overview.md
- Online Quantization: user_guide/quantization/online.md
- Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
- FP8 W8A8: user_guide/quantization/fp8.md
- Int8 W8A8: user_guide/quantization/int8.md
- BitsAndBytes W4: user_guide/quantization/bitsandbytes.md
- MXFP8 W8A8: user_guide/quantization/mxfp8.md
- MXFP4 W4A4: user_guide/quantization/mxfp4.md
- ModelOpt: user_guide/quantization/modelopt.md
- GGUF: user_guide/quantization/gguf.md
- AutoRound: user_guide/quantization/autoround.md
- msModelSlim: user_guide/quantization/msmodelslim.md
- ComfyUI: features/comfyui.md
- Fault Tolerance:
- Failure Modes: user_guide/fault_injection_reliability_matrix.md
@@ -114,7 +118,6 @@ nav:
- Shared Memory Connector: design/feature/omni_connectors/shared_memory_connector.md
- Yuanrong Store Connector: design/feature/omni_connectors/yuanrong_connector.md
- Yuanrong Transfer Engine Connector: design/feature/omni_connectors/yuanrong_transfer_engine_connector.md
- Quantization: design/feature/quantization.md
- Diffusion Acceleration:
- Parallelism:
- CFG-Parallel: design/feature/cfg_parallel.md
@@ -124,8 +127,9 @@ nav:
- Sequence Parallel: design/feature/sequence_parallel.md
- Tensor Parallel: design/feature/tensor_parallel.md
- VAE Patch Parallelism: design/feature/vae_parallel.md
- Attention Optimization:
- Attention Backends:
- Skip-Softmax: design/feature/skip_softmax.md
- Quantization: design/feature/quantization.md
- Cache-DiT: design/feature/cache_dit.md
- TeaCache: design/feature/teacache.md
- Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md
Binary file not shown.

Before

Width:  |  Height:  |  Size: 260 KiB

After

Width:  |  Height:  |  Size: 274 KiB

-41
View File
@@ -76,47 +76,6 @@ Please refer to the [test instructions](./ci/test_execution_guide.md) for detail
!!! warning
Currently, not all unit tests pass when run on CPU platforms. If you don't have access to a GPU platform to run unit tests locally, rely on the continuous integration system to run the tests for now.
### Using repository skills with coding agents
vLLM-Omni maintains [repository-scale skills](https://github.com/vllm-project/vllm-omni/tree/main/.claude/skills) for common coding, testing, performance, and review workflows. These skills capture repository-specific structure, validation requirements, and evidence standards that a general coding agent may not know.
If your coding agent discovers repository skills automatically, ask it to use the relevant skill by name. Otherwise, point it to the linked `SKILL.md` and ask it to read the file before changing code. Skills can be combined: use a domain skill for the implementation, `vllm-omni-test` for test and CI decisions, and `precheck-pr` before opening the pull request.
| Task | Repository skill |
| --- | --- |
| Add or extend a diffusion model | [`add-diffusion-model`](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/add-diffusion-model/SKILL.md) |
| Add or extend a text-to-speech model | [`add-tts-model`](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/add-tts-model/SKILL.md) |
| Add, select, or debug quantization | [`quantization`](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/quantization/SKILL.md) |
| Diagnose or optimize diffusion performance | [`diffusion-perf-opt`](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/diffusion-perf-opt/SKILL.md) |
| Upgrade an NPU model runner | [`vllm-omni-npu-model-runner-upgrade`](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/vllm-omni-npu-upgrade/SKILL.md) |
| Add a regression test, choose L1-L4 coverage, or wire Buildkite | [`vllm-omni-test`](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/vllm-omni-test/SKILL.md) |
| Self-check a branch before opening a pull request | [`precheck-pr`](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) |
| Perform a maintainer or reviewer review of an existing pull request or local branch | [`review-pr`](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/review-pr/SKILL.md) |
Use the following workflow when collaborating with a coding agent:
1. **Ground the task.** Provide the issue, RFC, or design document that defines the expected behavior. Ask the agent to state its assumptions and identify the affected module or feature contracts before editing code.
2. **Select the narrowest applicable skills.** Start with the domain skill that owns the change. Add `vllm-omni-test` whenever production behavior or tests change. For a bug fix, ask for the smallest reproducer and a regression test tied to the original symptom.
3. **Require executable evidence.** Ask the agent to run the narrowest relevant checks and report the exact commands and results. It must identify tests it could not run because of unavailable hardware, model weights, credentials, or dependencies; an unrun check is not a pass.
4. **Inspect the resulting diff.** Confirm that the implementation stays within the issue or RFC, does not overwrite unrelated work, and includes the required tests, documentation, and benchmark or accuracy evidence.
5. **Run the pre-submit workflow.** Use `precheck-pr`, then run the applicable local tests and pre-commit hooks before opening the pull request.
For example:
```text
Use add-diffusion-model and vllm-omni-test to implement the linked issue.
Read the relevant design documents first, keep the change within the issue,
and report every test command, result, and hardware validation gap.
```
```text
Use vllm-omni-test to turn this bug reproducer into the smallest deterministic
regression test. Prefer L1 CPU coverage; if the failure requires real weights
or serving, explain the required L2/L3 environment and provide the exact command.
```
Repository skills guide the agent, but they do not replace contributor judgment, the accepted issue or RFC, required test evidence, or maintainer review. Review all generated changes before submitting them.
## Issues
If you encounter a bug or have a feature request, please search existing issues first to see if it has already been reported. If not, please file a new issue, providing as much relevant information as possible.
+5 -13
View File
@@ -11,10 +11,6 @@ in the active navigation.
## Feature Design Documents
For user-facing configuration and current compatibility, see the
[Features overview](../features/README.md). A design document defines an
implementation contract; it is not, by itself, a general support claim.
### Runtime and stage execution
- [Disaggregated Inference](feature/disaggregated_inference.md)
@@ -34,10 +30,6 @@ implementation contract; it is not, by itself, a general support claim.
- [Yuanrong Store Connector](feature/omni_connectors/yuanrong_connector.md)
- [Yuanrong Transfer Engine Connector](feature/omni_connectors/yuanrong_transfer_engine_connector.md)
### Quantization
- [Quantization](feature/quantization.md)
### Diffusion acceleration
#### Parallelism
@@ -50,14 +42,14 @@ implementation contract; it is not, by itself, a general support claim.
- [Tensor Parallel](feature/tensor_parallel.md)
- [VAE Patch Parallelism](feature/vae_parallel.md)
#### Attention optimization
The [Diffusion Attention Backends](../user_guide/diffusion/attention_backends.md)
guide lists all selectable backends and their platform defaults. Skip-Softmax
is the backend-related optimization with a standalone feature-design contract:
#### Attention Backends
- [Skip-Softmax](feature/skip_softmax.md)
#### Quantization
- [Quantization](feature/quantization.md)
- [Cache-DiT](feature/cache_dit.md)
- [TeaCache](feature/teacache.md)
- [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
-1
View File
@@ -7,7 +7,6 @@ owners:
- "@yinpeiqi"
- "@fake0fan"
- "@Sy0307"
- "@Gaohan123"
primary_code_paths:
- vllm_omni/core/**
- vllm_omni/worker/**
-1
View File
@@ -4,7 +4,6 @@ kind: module
status: draft
owners:
- "@alex-jw-brooks"
- "@Bounty-hunter"
primary_code_paths:
- vllm_omni/benchmarks/**
- vllm_omni/entrypoints/cli/benchmark/**
-1
View File
@@ -5,7 +5,6 @@ status: draft
owners:
- "@Isotr0py"
- "@princepride"
- "@SamitHuang"
primary_code_paths:
- vllm_omni/diffusion/cache/**
related_code_paths:
@@ -3,10 +3,9 @@ title: Diffusion Model Integration
kind: module
status: draft
owners:
- "@fhfuih"
- "@Bounty-hunter"
- "@Isotr0py"
- "@wtomin"
- "@RuixiangMa"
- "@ZJY0516"
primary_code_paths:
- vllm_omni/diffusion/models/**
- vllm_omni/diffusion/model_loader/**
@@ -6,7 +6,6 @@ owners:
- "@Isotr0py"
- "@princepride"
- "@SamitHuang"
- "@fhfuih"
primary_code_paths:
- vllm_omni/diffusion/sched/**
- vllm_omni/diffusion/executor/**
+1 -2
View File
@@ -4,8 +4,7 @@ kind: module
status: draft
owners:
- "@yuanheng-zhao"
- "@david6666666"
- "@lishunyang12"
- "@Isotr0py"
primary_code_paths:
- vllm_omni/diffusion/offloader/**
related_code_paths:
@@ -6,7 +6,6 @@ owners:
- "@princepride"
- "@yuanheng-zhao"
- "@xuechendi"
- "@Bounty-hunter"
primary_code_paths:
- vllm_omni/diffusion/distributed/**
- vllm_omni/diffusion/attention/parallel/**
-1
View File
@@ -6,7 +6,6 @@ architecture_state: current-plus-in-flight
owners:
- "@alex-jw-brooks"
- "@linyueqian"
- "@NickCao"
document_stewards:
- "@hsliuustc0106"
- "@Gaohan123"
+1 -8
View File
@@ -8,15 +8,8 @@ owners:
- "@xuechendi"
primary_code_paths:
- vllm_omni/platforms/**
related_code_paths:
- vllm_omni/attention/**
primary_code_path_owners:
# Additional owners scoped to attention kernels (quantized/custom attention)
- paths:
- vllm_omni/attention/**
owners:
- "@david6666666"
- "@lishunyang12"
related_code_paths: []
depends_on:
- ar_runtime.md
- diffusion/index.md
@@ -6,8 +6,6 @@ architecture_state: current
owners:
- "@Sy0307"
- "@amy-why-3459"
- "@Gaohan123"
- "@alex-jw-brooks"
document_stewards:
- "@hsliuustc0106"
- "@Gaohan123"
-8
View File
@@ -9,14 +9,6 @@ primary_code_paths:
- vllm_omni/model_executor/**
- vllm_omni/model_extras/**
- vllm_omni/plugins/**
primary_code_path_owners:
# Additional owners scoped to a sub-path of model_executor/**
- paths:
- vllm_omni/model_executor/models/**
owners:
- "@Sy0307"
- "@amy-why-3459"
- "@linyueqian"
related_code_paths:
- vllm_omni/transformers_utils/**
- vllm_omni/tokenizers/**
-1
View File
@@ -4,7 +4,6 @@ kind: module
status: draft
owners:
- "@lishunyang12"
- "@vraiti"
primary_code_paths:
- vllm_omni/metrics/**
related_code_paths:
-67
View File
@@ -1,67 +0,0 @@
# Features
Use this section to choose and configure vLLM-Omni runtime and optimization
features. The navigation follows the active
[Feature Design](../design/index.md#feature-design-documents) taxonomy whenever
an implementation contract has a user-facing workflow.
!!! note "User guides, design documents, and recipes"
User guides explain how to enable a feature and where it is supported.
Design documents define internal contracts and do not, by themselves,
imply general or production support. For model-specific launch flags and
hardware requirements, use the linked examples and validated recipes from
[Supported Models](../models/supported_models.md).
## Runtime and Stage Execution
| Goal | User guide | Related design contract |
| --- | --- | --- |
| Choose serial, batched, step-wise, or streaming diffusion execution | [Execution Modes and Streaming](../user_guide/diffusion/execution_modes.md) | [Diffusion Continuous Batching](../design/feature/diffusion_continuous_batching.md), [Async Diffusion Output](../design/feature/async_diffusion_output.md) |
| Reclaim stage memory without restarting the server | [Sleep Mode](sleep_mode.md) | Runtime lifecycle behavior is documented in the user guide |
Some runtime designs are deliberately not promoted as standalone User Guide
features yet:
- [Disaggregated Inference](../design/feature/disaggregated_inference.md) and
[OmniConnector implementations](../design/index.md#communication) describe
topology and transport contracts. Practical deployment configuration remains
under [Pipeline and deploy configurations](../configuration/stage_configs.md).
- [Async Chunk](../design/feature/async_chunk.md) and
[Async Omni Output Materialization](../design/feature/omni_async_output_materialization.md)
are model- and pipeline-dependent. Use the selected model's deploy
configuration and recipe for the supported settings.
- [Automatic Prefix Caching](../design/feature/prefix_caching.md) remains a
design-level contract until its user-facing configuration and compatibility
surface is consolidated.
## Quantization
Quantization is a cross-model feature rather than a diffusion-only
optimization. The unified [`quantization_config` guide](../user_guide/quantization/overview.md)
covers diffusion-only models, multi-stage omni/TTS models, and multi-stage
diffusion models. Its [design contract](../design/feature/quantization.md)
defines the shared configuration and backend extension points.
## Diffusion Acceleration
| Goal | User guide | Related design contract |
| --- | --- | --- |
| Compare acceleration methods and supported combinations | [Overview](../user_guide/diffusion_features.md), [Feature Compatibility](../user_guide/feature_compatibility.md) | [Diffusion acceleration designs](../design/index.md#diffusion-acceleration) |
| Move weights between host and device memory | [CPU Offloading](../user_guide/diffusion/cpu_offload.md) | [Distributed Layerwise Offload](../design/feature/distributed_layerwise_offload.md) |
| Reuse denoising computation | [Cache Acceleration](../user_guide/diffusion/cache_acceleration/cache_dit.md) | [Cache-DiT](../design/feature/cache_dit.md), [TeaCache](../design/feature/teacache.md) |
| Distribute diffusion work across devices | [Parallelism](../user_guide/diffusion/parallelism/overview.md) | [Parallelism designs](../design/index.md#parallelism) |
| Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Skip-Softmax](../design/feature/skip_softmax.md) |
| Compile repeated diffusion regions | [Regional Compilation](../user_guide/diffusion/regional_compilation.md) | User-facing optimization guide |
| Add generated video frames | [Frame Interpolation](../user_guide/diffusion/frame_interpolation.md) | User-facing extension guide |
| Reduce diffusion model startup time | [Startup and Loading](../user_guide/diffusion/startup_and_loading.md) | User-facing loading guide |
| Apply diffusion adapters | [LoRA](../user_guide/diffusion/lora.md) | User-facing extension guide |
The [Pipeline Parallelism guide](../user_guide/diffusion/parallelism/pipeline_parallel.md)
remains available by direct link, but it is not promoted in the primary User
Guide navigation until its user-facing support and placement are settled.
## Experimental
[Session State Manager](session_state_manager.md) is opt-in and experimental.
Its APIs and compatibility may change without notice.
-1
View File
@@ -123,7 +123,6 @@ def extract_omni_serve_subparser_init():
"TrackingNamespace": None, # type annotation only
"make_arg_parser": lambda parser: parser, # no-op for doc
"_ensure_vllm_platform": lambda: None, # no-op for doc
"_nonneg_finite_float": float,
"VLLM_SUBCMD_PARSER_EPILOG": "",
"logger": logger,
"DummySubparsers": DummySubparsers,
+12 -37
View File
@@ -24,23 +24,6 @@ SERVING_MODE_TITLES = {
"online_serving": "Online serving",
}
# Keep the Examples navigation focused on reusable modality/task entry points.
# Model-specific scripts remain available under ``examples/`` and are surfaced
# from these shared pages instead of getting one page per model.
GENERAL_EXAMPLE_SLUGS = frozenset(
{
"image_to_image",
"image_to_video",
"speech_to_video",
"text_to_audio",
"text_to_image",
"text_to_speech",
"text_to_video",
"x_to_text",
"x_to_video_audio",
}
)
def load_model_display_names() -> dict[str, str]:
try:
@@ -63,13 +46,6 @@ def load_model_display_names() -> dict[str, str]:
MODEL_DISPLAY_NAMES = load_model_display_names()
def is_general_example(example: "Example") -> bool:
"""Return whether an example should get a generated serving page."""
if example.category not in SERVING_MODE_TITLES:
return True
return example.path.stem in GENERAL_EXAMPLE_SLUGS
def fix_case(text: str) -> str:
subs = {
"api": "API",
@@ -364,21 +340,22 @@ def update_nav_file(examples: list[Example]):
examples_by_category[category] = []
examples_by_category[category].append(example)
def category_items(category: str) -> list[dict[str, str]]:
items = []
for example in sorted(examples_by_category[category], key=lambda e: e.path.stem):
# Build the new Examples section - start with preserved items
examples_section = preserved_items.copy()
# Add examples grouped by category, sorted by category name
for category in sorted(examples_by_category.keys()):
category_examples = sorted(examples_by_category[category], key=lambda e: e.path.stem)
category_items = []
for example in category_examples:
doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
rel_path = doc_path.relative_to(ROOT_DIR / "docs")
items.append({example.nav_title: rel_path.as_posix()})
return items
category_items.append({example.nav_title: rel_path.as_posix()})
# Build the new Examples section - start with preserved items.
examples_section = preserved_items.copy()
for category in sorted(examples_by_category.keys()):
items = category_items(category)
if items:
if category_items:
# Format category name (e.g., "offline_inference" -> "Offline Inference")
category_title = fix_case(category.replace("_", " ").title())
examples_section.append({category_title: items})
examples_section.append({category_title: category_items})
# Update the nav structure
nav_list[user_guide_idx]["User Guide"][examples_idx]["Examples"] = examples_section
@@ -413,8 +390,6 @@ def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
for path in category.glob("*/*.md"):
examples.append(Example(path.parent, category.stem))
examples = [example for example in examples if is_general_example(example)]
# Generate the example documentation
for example in sorted(examples, key=lambda e: e.path.stem):
example_name = f"{example.path.stem}.md"
+78 -88
View File
@@ -17,94 +17,84 @@ th {
}
</style>
| Architecture | Models | Example HF Models | NVIDIA GPU | AMD GPU | Ascend NPU | Intel GPU | Recipe |
|--------------|--------|-------------------|------------|---------|-----|-----------|--------|
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md) |
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | — |
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ | — |
| `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) |
| `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` | ✅︎ | | | | — |
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `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` | ✅︎ | | | | — |
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | — |
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
| `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` | ✅︎ | | | | — |
| `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` | ✅︎ | | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | — |
| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | — |
| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | — |
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — |
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — |
| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — |
| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — |
| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — |
| `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` | ✅︎ | | | | — |
| `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` | ✅︎ | | | | — |
| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | | — |
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | | — |
| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | | [Published](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b) |
| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | | — |
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ | — |
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | | — |
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | | — |
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
| `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` | ✅︎ | | | | — |
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
The `Recipe` column links the corresponding deployment recipe. `Published` is
preferred when a corresponding `recipes.vllm.ai` page exists; otherwise the
link points to the repository recipe. `—` means that no recipe link has been
added for that model row yet.
For rows with an audited recipe link, the hardware checkmarks reflect hardware
explicitly documented in that recipe. Other rows retain the implementation
support metadata until their recipe is audited.
| Architecture | Models | Example HF Models | NVIDIA GPU | AMD GPU | Ascend NPU | Intel GPU |
|--------------|--------|-------------------|------------|---------|-----|-----------|
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | |
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `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` | ✅︎ | | | |
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `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` | ✅︎ | | | |
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | |
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | | | |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-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` | ✅︎ | | | |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | |
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ |
| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ |
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ |
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ |
| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | |
| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | |
| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `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` | ✅︎ | | | |
| `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` | ✅︎ | | | |
| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | |
| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | |
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | |
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ |
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | |
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
| `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` | ✅︎ | | | |
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | |
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
+15 -12
View File
@@ -283,20 +283,23 @@ blocks — use Ulysses SP (`ring_degree=1`).
### Which geometries run sparse
MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when
necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2
does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the
promoted rows remain visible to every query and key.
Sparsity is only applied when the video segment is a **multiple of 128 rows**, where the row count
is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own
tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows
that selection may then drop. Any resolution still runs — an unaligned geometry falls back to dense
attention and logs `RAINFUSION_ATTN staying dense` with the row count it computed — but it gets no
speedup, so pick an aligned geometry when you want one.
For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row
128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix.
MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask
and kernel tiling aligned, at the cost of retaining more blocks.
Alignment is necessary, not sufficient, for good quality. `rf_v2` groups video positions into 8x8
spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact
patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns
are peeled off and appended as a flat run instead of being tiled, so they get pooled with spatially
distant positions and selection can no longer rank them meaningfully. This is silent, and invisible
to a `sparsity=0` check, because a fully populated mask does not care how blocks are grouped.
8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8
creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent
grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release
that implements protected video tails for rf_v2.
"Multiple of 8" on the latent grid means **width and height that are multiples of 256**. Aligned
resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`,
which you can buy back with `start_step`.
## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
@@ -164,12 +164,10 @@ In `DiffusionParallelConfig`:
| `ulysses_degree` | int | 1 | Number of GPUs for Ulysses-SP. Uses all-to-all communication. |
| `ring_degree` | int | 1 | Number of GPUs for Ring-Attention. Uses P2P ring communication. |
| `ulysses_mode` | str | `"default"` | Ulysses attention mode. Set to `"advanced_uaa"` to handle arbitrary sequence lengths and head counts without padding. |
| `mask_sp_padding` | bool | `False` | When the sequence length is not divisible by the SP size, tokens are auto-padded with zeros. Set to `True` to mask those padding tokens (strict, but uses the slower varlen attention path); the default `False` leaves them unmasked, keeping the fast path with negligible numerical impact. |
**Notes:**
- Total sequence parallel size equals to `ulysses_degree × ring_degree`
- Degrees must evenly divide the sequence length for optimal performance (or use `ulysses_mode="advanced_uaa"` for Ulysses-SP)
- `mask_sp_padding` is an experimental feature, currently only supported by `Wan2.2`, `Wan2.2 Vace`, `Qwen-Image`, `Flux 2`, and `HunyuanVideo 1.5`
## Best Practices
@@ -1,60 +0,0 @@
# Diffusion Startup and Loading
Large diffusion models can take several minutes to load at startup. vLLM-Omni
loads safetensors shards in parallel to reduce this initialization time.
Multi-thread weight loading is enabled by default with four threads. No
configuration is needed for the default behavior.
## Configuration
| Parameter | CLI flag | Default | Description |
| --- | --- | --- | --- |
| `enable_multithread_weight_load` | `--disable-multithread-weight-load` | `True` | Pass the flag to disable multi-thread loading |
| `num_weight_load_threads` | `--num-weight-load-threads` | `4` | Number of parallel weight-loading threads |
!!! tip
The default balances startup speed and disk I/O contention. Fast NVMe
storage may benefit from more threads, while network storage or hard disks
may not.
## Online Serving
```bash
# Default: multi-thread loading with four threads
vllm serve Qwen/Qwen-Image --omni --port 8091
# Increase the thread count
vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni \
--num-weight-load-threads 8
# Disable multi-thread loading
vllm serve Qwen/Qwen-Image --omni --disable-multithread-weight-load
```
## Offline Inference
```python
from vllm_omni import Omni
# Default: multi-thread loading with four threads
omni = Omni(model="Qwen/Qwen-Image")
# Increase the thread count
omni = Omni(
model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
num_weight_load_threads=8,
)
```
## Reference Benchmarks
The following measurements were collected on NVIDIA H800 hardware. Treat them
as reference results rather than a guarantee for other storage or hardware
configurations.
| Model | Sequential loading | Multi-thread loading | Speedup |
| --- | ---: | ---: | ---: |
| **Qwen/Qwen-Image** (53.7 GiB) | 168 s | 27 s | **6.2x** |
| **Wan-AI/Wan2.2-I2V-A14B-Diffusers** (64.5 GiB) | 283 s | 56 s | **5.1x** |
+86 -39
View File
@@ -49,7 +49,7 @@ Parallelism methods distribute computation across GPUs without quality loss (mat
| Method | Description | Best For |
|--------|-------------|----------|
| **[Multi-Thread Weight Loading](diffusion/startup_and_loading.md)** | Loads safetensors shards in parallel using a thread pool | All diffusion models; reduces startup from minutes to seconds |
| **[Multi-Thread Weight Loading](#multi-thread-weight-loading)** | Loads safetensors shards in parallel using a thread pool | All diffusion models; reduces startup from minutes to seconds |
**Note:** Some acceleration methods can be combined together for optimized performance. See [Feature Compatibility Table](#feature-compatibility) and [Feature Compatibility Tutorial](feature_compatibility.md) for detailed configuration examples.
@@ -100,9 +100,7 @@ The following tables show which models support each feature:
- **🔀SP (Ulysses & Ring)**: Includes both Ulysses-SP and Ring-Attention methods
- ✅ = Fully supported
- ✅* = Supported with the constraint listed below the table
- ❌ = Not supported
- ❓ = Not verified; not recommended
> Notes:
@@ -121,7 +119,7 @@ The following tables show which models support each feature:
| **FLUX.2-dev** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **Hidream-I1-Full** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅* |
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
| **Krea 2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
@@ -147,7 +145,6 @@ The following tables show which models support each feature:
> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload.
> 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.
> 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
### VideoGen
@@ -158,15 +155,11 @@ The following tables show which models support each feature:
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* |
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **LongCat-Video-Avatar-1.5** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **MiniMax-H3** | ✅ (FL2VA) | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ |
> **Step execution note:** Helios supports single-request step execution only;
> use `max_num_seqs=1`.
| **MiniMax-H3** | ❌ | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ |
**Frame Interpolation Support**
@@ -196,7 +189,7 @@ The following tables show which models support each feature:
| **🔀Ring-Attn** | ✅ | ✅ | ✅ | | | | | | | | | | | |
| **🔀CFG-Parallel** | ✅ | ✅ | ✅ | ✅ | | | | | | | | | | |
| **🔀Tensor Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | | |
| **🔀HSDP** | ❓ | ❓ | ✅ | ❓ | ✅ | ❌ | | | | | | | | |
| **🔀HSDP** | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | | | | | | | | |
| **🔀Expert Parallel** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | | | | | | |
| **💾CPU Offloading (Layerwise)** | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | | | | | |
| **💾CPU Offloading (Module-wise)** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ❌ | | | | | |
@@ -207,43 +200,97 @@ The following tables show which models support each feature:
!!! info
1. HSDP can be combined with Ulysses-SP or CFG-Parallel. Tensor Parallel
and HSDP are not compatible; other HSDP combinations in the table
remain unverified.
1. Tensor Parallel and HSDP are not compatible.
2. TeaCache and Cache-DiT are not compatible.
3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
4. The CPU Offloading (Layerwise) row describes local layerwise offload.
Multi-device Distributed Layerwise Offload has a separate topology and
compatibility matrix in the [CPU Offloading guide](diffusion/cpu_offload.md#distributed-layerwise-offloading).
5. The compatibility matrix uses FP8 as the representative quantization method.
4. CPU Offloading (Layerwise) supports single-card for now.
5. Using FP8-Quant as an example of qunatization methods.
6. Step Execution is not compatible with any diffusion cache backend. LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).
## Multi-Thread Weight Loading
The loading guide now lives at [Diffusion Startup and
Loading](diffusion/startup_and_loading.md). This heading remains so existing
links to this section continue to work.
Large diffusion models can take several minutes to load weights at startup (e.g., ~3 min for Qwen-Image, ~5 min for Wan2.2 I2V 14B). Multi-thread weight loading speeds up this process by loading safetensors shards in parallel using a thread pool instead of sequentially.
This optimization is **enabled by default** with 4 threads. No configuration is needed for the default behavior.
### Configuration
| Parameter | CLI Flag | Default | Description |
|-----------|----------|---------|-------------|
| `enable_multithread_weight_load` | `--disable-multithread-weight-load` | `True` (enabled) | Pass the flag to disable multi-thread loading |
| `num_weight_load_threads` | `--num-weight-load-threads` | `4` | Number of threads for parallel weight loading |
!!! tip
The default of 4 threads balances speed and disk I/O contention. On fast NVMe storage you may benefit from more threads (e.g., 8). On HDD or network storage, the default of 4 avoids saturating I/O bandwidth.
### Online Serving
```bash
# Default (multi-thread enabled, 4 threads)
vllm serve Qwen/Qwen-Image --omni --port 8091
# Custom thread count
vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --num-weight-load-threads 8
# Disable multi-thread loading
vllm serve Qwen/Qwen-Image --omni --disable-multithread-weight-load
```
### Offline Inference
```python
from vllm_omni import Omni
# Default (multi-thread enabled, 4 threads)
omni = Omni(model="Qwen/Qwen-Image")
# Custom thread count
omni = Omni(
model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
num_weight_load_threads=8,
)
```
### Benchmarks
Measured on NVIDIA H800:
| Model | Before | After | Speedup |
|-------|--------|-------|---------|
| **Qwen/Qwen-Image** (53.7 GiB) | 168s | 27s | **6.2x** |
| **Wan-AI/Wan2.2-I2V-A14B-Diffusers** (64.5 GiB) | 283s | 56s | **5.1x** |
## Learn More
The Diffusion Acceleration navigation groups the remaining guides as follows:
**Cache Acceleration:**
| Area | Guide |
| --- | --- |
| Compatibility | [Feature Compatibility](feature_compatibility.md) |
| CPU offloading | [CPU Offloading](diffusion/cpu_offload.md) |
| Cache acceleration | [TeaCache](diffusion/cache_acceleration/teacache.md), [Cache-DiT](diffusion/cache_acceleration/cache_dit.md) |
| Parallelism | [Parallelism Overview](diffusion/parallelism/overview.md) |
| Attention | [Attention Backends](diffusion/attention_backends.md) |
| Compilation | [Regional Compilation](diffusion/regional_compilation.md) |
| Video extension | [Frame Interpolation](diffusion/frame_interpolation.md) |
| Startup | [Startup and Loading](diffusion/startup_and_loading.md) |
| Adapters | [LoRA](diffusion/lora.md) |
- **[TeaCache Configuration Guide](diffusion/cache_acceleration/teacache.md)** - Parameter tuning, performance tips, troubleshooting
- **[Cache-DiT Advanced Guide](diffusion/cache_acceleration/cache_dit.md)** - DBCache, TaylorSeer, SCM techniques and optimization
Related cross-model and runtime features are documented separately:
**Parallelism Methods:**
- [Quantization](quantization/overview.md) covers diffusion-only models,
multi-stage omni/TTS models, and multi-stage diffusion models.
- [Execution Modes and Streaming](diffusion/execution_modes.md) covers the
diffusion runtime, including batching, step execution, and streaming output.
- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, Pipeline Parallelism, HSDP, and Expert Parallelism
**Memory Optimization:**
- **[CPU Offload Guide](diffusion/cpu_offload.md)** - Offload model components to CPU, reduce GPU memory usage
- **[VAE Parallelism Guide](diffusion/parallelism/vae_parallelism.md)** - Distribute VAE decode work across GPUs for high-resolution images and videos
- **[Quantization](quantization/overview.md)** - Quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models
**Extensions:**
- **[LoRA Inference Guide](diffusion/lora.md)** - Low-Rank Adaptation for style customization and fine-tuning
- **[Frame Interpolation Guide](diffusion/frame_interpolation.md)** - Worker-side post-generation video frame interpolation for smoother motion
**Execution Modes:**
- **[Diffusion Execution Modes](diffusion/execution_modes.md)** - Configure request batching, step execution, continuous batching, and streaming output
**Startup Optimization:**
- **[Multi-Thread Weight Loading](#multi-thread-weight-loading)** - Speed up model startup by loading safetensors shards in parallel
**Advanced Topics:**
- **[Feature Compatibility](feature_compatibility.md)** - How to combine multiple features for maximum performance
@@ -99,7 +99,7 @@ a native vLLM-Omni pipeline for those cases.
### Attention Backends
The diffusers backend converts
[vLLM-Omni standard of attention backend setting](../../diffusion/attention_backends.md)
[vLLM-Omni standard of attention backend setting](../../../docs/user_guide/diffusion/attention_backends.md)
to [diffusers standard](https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends).
Specifically for `FLASH_ATTN`, it will first attempt to use FlashAttention-3 and then FlashAttention-2.
@@ -33,25 +33,6 @@ This folder provides a unified CLI script for image-to-video generation using vL
Default model: `Wan-AI/Wan2.2-I2V-A14B-Diffusers`.
### LingBot-Video TI2V
LingBot-Video accepts exactly one first-frame image and uses the same dense or
MoE checkpoint as its T2I and T2V modes:
```bash
python examples/offline_inference/image_to_video/image_to_video.py \
--model robbyant/lingbot-video-dense-1.3b \
--model-class-name LingBotVideoPipeline \
--image /path/to/input.png \
--prompt "the fox looks toward the camera" \
--height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
--guidance-scale 3.0 --flow-shift 3.0 --fps 24 \
--output lingbot_ti2v.mp4
```
The input image keeps its original geometry until the LingBot pipeline applies
its model-specific resize and center crop.
## Prerequisites
Download the example image used in the snippets below:
@@ -1,319 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import json
import re
import subprocess
import tempfile
import wave
from functools import lru_cache
from pathlib import Path
from typing import Any
import imageio
import numpy as np
import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.models.longcat_video.pipeline_longcat_video_avatar import (
prepare_longcat_video_avatar_model_for_omni,
)
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
def parse_args():
parser = argparse.ArgumentParser(description="Generate LongCat-Video-Avatar A2V/AI2V video.")
parser.add_argument("--model", default="meituan-longcat/LongCat-Video-Avatar-1.5")
parser.add_argument("--base-model-dir", default=None, help="Optional local LongCat-Video base model directory.")
parser.add_argument("--stage", choices=["at2v", "ai2v"], default="at2v")
parser.add_argument("--prompt", default="A person speaks calmly while facing the camera.")
parser.add_argument("--negative-prompt", default="low quality, blurry, watermark, text")
parser.add_argument("--audio", default=None, help="Input speech audio path.")
parser.add_argument("--image", default=None, help="Input reference image path for AI2V.")
parser.add_argument("--input-json", default=None, help="Official LongCat Avatar JSON case for AI2V.")
parser.add_argument("--output", default="longcat_avatar_output.mp4")
parser.add_argument("--resolution", choices=["480p", "720p"], default="480p")
parser.add_argument("--num-frames", type=int, default=93)
parser.add_argument("--num-segments", default="1", help="Number of AVC segments, or 'auto' to cover full audio.")
parser.add_argument("--num-cond-frames", type=int, default=13)
parser.add_argument("--ref-img-index", type=int, default=10)
parser.add_argument("--mask-frame-range", type=int, default=3)
parser.add_argument("--fps", type=int, default=25)
parser.add_argument("--num-inference-steps", type=int, default=8)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--use-distill", action=argparse.BooleanOptionalAction, default=True)
parser.add_argument("--use-int8", action=argparse.BooleanOptionalAction, default=True)
parser.add_argument(
"--build-components-on-gpu",
action="store_true",
default=False,
help="Build large Avatar components directly on GPU for faster startup. Requires more peak VRAM.",
)
parser.add_argument("--use-kv-cache", action=argparse.BooleanOptionalAction, default=True)
parser.add_argument("--offload-kv-cache", action=argparse.BooleanOptionalAction, default=False)
parser.add_argument("--enforce-eager", action="store_true")
return parser.parse_args()
def _extract_frames(output: Any):
if isinstance(output, list):
output = output[0] if output else None
if isinstance(output, OmniRequestOutput):
if not output.images:
raise ValueError("No video frames found in OmniRequestOutput.")
frames = output.images[0] if len(output.images) == 1 else output.images
if isinstance(frames, tuple) and len(frames) == 2:
frames = frames[0]
if isinstance(frames, dict):
frames = frames.get("frames") or frames.get("video")
else:
frames = output
if frames is None:
raise ValueError("No video frames found in output.")
return frames
def _frames_to_uint8(frames):
frames_np = np.asarray([np.asarray(frame) for frame in frames])
if frames_np.ndim == 5:
frames_np = frames_np[0]
if np.issubdtype(frames_np.dtype, np.floating):
frames_np = (np.clip(frames_np, 0.0, 1.0) * 255).round()
return np.clip(frames_np, 0, 255).astype(np.uint8)
@lru_cache(maxsize=1)
def _ffmpeg_exe() -> str:
try:
import imageio_ffmpeg
return imageio_ffmpeg.get_ffmpeg_exe()
except Exception as exc:
raise RuntimeError("LongCat Avatar example requires the Python imageio[ffmpeg] dependency.") from exc
def _wav_duration(path: Path) -> float:
with wave.open(str(path), "rb") as audio:
return audio.getnframes() / float(audio.getframerate())
def _save_video_with_audio(frames, output_path: Path, audio_path: str, fps: int) -> None:
frames_np = _frames_to_uint8(frames)
output_path.parent.mkdir(parents=True, exist_ok=True)
with tempfile.TemporaryDirectory(prefix="longcat_avatar_export_") as tmp:
tmp_dir = Path(tmp)
temp_video = tmp_dir / "video.mp4"
crop_audio = tmp_dir / "audio.wav"
crop_video = tmp_dir / "video_cropped.mp4"
writer = imageio.get_writer(str(temp_video), fps=fps, quality=5)
try:
for frame in frames_np:
writer.append_data(frame)
finally:
writer.close()
duration = len(frames_np) / fps
ffmpeg = _ffmpeg_exe()
subprocess.run([ffmpeg, "-y", "-i", audio_path, "-t", f"{duration}", str(crop_audio)], check=True)
audio_duration = _wav_duration(crop_audio)
subprocess.run(
[
ffmpeg,
"-y",
"-i",
str(temp_video),
"-t",
f"{audio_duration}",
"-c:v",
"copy",
str(crop_video),
],
check=True,
)
subprocess.run(
[
ffmpeg,
"-y",
"-i",
str(crop_video),
"-i",
str(crop_audio),
"-c:v",
"libx264",
"-c:a",
"aac",
"-shortest",
str(output_path),
],
check=True,
)
def _infer_asset_root_from_input_json(input_json: Path) -> Path:
resolved = input_json.resolve()
parts = resolved.parts
if "assets" not in parts:
return resolved.parent
asset_idx = parts.index("assets")
return Path(*parts[:asset_idx])
def _resolve_case_path(path: str, asset_root: Path) -> str:
raw_path = Path(path)
if raw_path.is_absolute():
return str(raw_path)
candidate = asset_root / raw_path
return str(candidate if candidate.exists() else raw_path)
def _speaker_sort_key(name: str) -> tuple[int, str]:
match = re.search(r"(\d+)$", str(name))
return (int(match.group(1)) if match else 10_000, str(name))
def _load_input_case(input_json: str) -> tuple[dict[str, Any], Path]:
input_json_path = Path(input_json)
with input_json_path.open("r", encoding="utf-8") as f:
case = json.load(f)
return case, _infer_asset_root_from_input_json(input_json_path)
def _create_merged_audio_for_case(
audio_paths: dict[str, str],
audio_type: str,
output_path: Path,
) -> str:
sorted_paths = [path for _, path in sorted(audio_paths.items(), key=lambda item: _speaker_sort_key(str(item[0])))]
if len(sorted_paths) == 1:
return sorted_paths[0]
inputs = [arg for path in sorted_paths for arg in ("-i", path)]
audio_type = audio_type.lower()
if audio_type == "para":
filter_complex = f"amix=inputs={len(sorted_paths)}:duration=longest:normalize=0"
elif audio_type == "add":
filter_complex = "".join(f"[{idx}:a]" for idx in range(len(sorted_paths)))
filter_complex += f"concat=n={len(sorted_paths)}:v=0:a=1"
else:
raise NotImplementedError(f"Unsupported LongCat Avatar audio_type {audio_type!r}.")
subprocess.run([_ffmpeg_exe(), "-y", *inputs, "-filter_complex", filter_complex, str(output_path)], check=True)
return str(output_path)
def main():
args = parse_args()
input_case = None
input_asset_root = None
audio_for_mux = args.audio
if args.input_json:
input_case, input_asset_root = _load_input_case(args.input_json)
args.stage = "ai2v"
args.prompt = input_case.get("prompt") or args.prompt
args.image = _resolve_case_path(input_case["cond_image"], input_asset_root)
args.audio = None
if not args.input_json and not args.audio:
raise ValueError("--audio is required unless --input-json is provided.")
if args.stage == "ai2v" and not args.image:
raise ValueError("--image is required when --stage=ai2v.")
model_for_omni = prepare_longcat_video_avatar_model_for_omni(args.model, args.use_int8)
additional_config = {
"model_type": "avatar-v1.5",
"resolution": args.resolution,
"use_distill": args.use_distill,
"use_int8": args.use_int8,
"build_components_on_gpu": args.build_components_on_gpu,
}
if args.base_model_dir:
additional_config["base_model_dir"] = args.base_model_dir
omni = Omni(
model=model_for_omni,
model_class_name="LongCatVideoAvatarPipeline",
dtype="bfloat16",
enforce_eager=args.enforce_eager,
additional_config=additional_config,
parallel_config=DiffusionParallelConfig(
ulysses_degree=1,
ring_degree=1,
cfg_parallel_size=1,
tensor_parallel_size=1,
vae_patch_parallel_size=1,
pipeline_parallel_size=1,
),
)
audio_input: str | dict[str, str] | None = args.audio
additional_information = {"stage": args.stage, "resolution": args.resolution}
if input_case is not None and input_asset_root is not None:
audio_input = {
speaker: _resolve_case_path(path, input_asset_root)
for speaker, path in sorted(
input_case["cond_audio"].items(),
key=lambda item: _speaker_sort_key(str(item[0])),
)
}
additional_information.update(
{
"input_json": args.input_json,
"audio_type": input_case.get("audio_type", "para"),
"bbox": input_case.get("bbox", {}),
}
)
prompt = {
"prompt": args.prompt,
"negative_prompt": args.negative_prompt,
"multi_modal_data": {"audio": audio_input},
"additional_information": additional_information,
}
if args.stage == "ai2v":
prompt["multi_modal_data"]["image"] = args.image
try:
output = omni.generate(
prompt,
OmniDiffusionSamplingParams(
generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed),
seed=args.seed,
guidance_scale=1.0 if args.use_distill else 4.0,
guidance_scale_2=1.0 if args.use_distill else 4.0,
num_inference_steps=args.num_inference_steps,
num_frames=args.num_frames,
fps=args.fps,
extra_args={
"stage": args.stage,
"resolution": args.resolution,
"save_fps": args.fps,
"use_distill": args.use_distill,
"use_int8": args.use_int8,
"input_json": args.input_json,
"num_segments": args.num_segments,
"num_cond_frames": args.num_cond_frames,
"ref_img_index": args.ref_img_index,
"mask_frame_range": args.mask_frame_range,
"use_kv_cache": args.use_kv_cache,
"offload_kv_cache": args.offload_kv_cache,
},
),
use_tqdm=True,
)
finally:
omni.close()
with tempfile.TemporaryDirectory(prefix="longcat_avatar_audio_mux_") as tmp:
if input_case is not None and isinstance(audio_input, dict):
audio_for_mux = _create_merged_audio_for_case(
audio_input,
input_case.get("audio_type", "para"),
Path(tmp) / "merged_audio.wav",
)
_save_video_with_audio(_extract_frames(output), Path(args.output), audio_for_mux, args.fps)
print(f"Saved video to {args.output}")
if __name__ == "__main__":
main()
@@ -43,23 +43,6 @@ This folder provides several entrypoints for experimenting with text-to-image di
Default model: `Qwen/Qwen-Image`
### LingBot-Video
LingBot-Video uses the same checkpoint for image and video generation. Run the
shared T2I entry point with:
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model robbyant/lingbot-video-dense-1.3b \
--prompt "a red fox standing in fresh snow" \
--height 192 --width 320 --num-inference-steps 2 \
--guidance-scale 3.0 --extra-body '{"flow_shift":3.0}' \
--output lingbot_t2i.png
```
The LingBot prompt builder selects the `image` output modality and requests one
frame. Passing an input image is not supported in LingBot T2I mode.
## Quick Start
### Python API
@@ -8,9 +8,7 @@ import time
from pathlib import Path
from typing import Any
import numpy as np
import torch
from diffusers.utils import numpy_to_pil
from vllm_omni.diffusion.data import logger
from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
@@ -56,17 +54,6 @@ def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]
parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
def _normalize_images_for_save(images: list[Any]) -> list[Any]:
"""Convert NumPy diffusion outputs to PIL images before saving."""
normalized = []
for image in images:
if isinstance(image, np.ndarray):
normalized.extend(numpy_to_pil(image))
else:
normalized.append(image)
return normalized
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
parser.add_argument(
@@ -632,7 +619,6 @@ def main():
if not images:
raise ValueError("No images found in request_output")
images = _normalize_images_for_save(images)
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
@@ -55,21 +55,6 @@ python text_to_video.py \
--output vace_t2v_output.mp4
```
### LingBot-Video
The dedicated LingBot example explicitly selects the video output modality:
```bash
python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
--model robbyant/lingbot-video-dense-1.3b \
--height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
--output lingbot_t2v.mp4
```
Use the same command with `robbyant/lingbot-video-moe-30b-a3b` for the MoE
checkpoint. Requested frame counts are rounded upward to the causal
VAE's `4n+1` grid.
### LTX-2
```bash
@@ -20,9 +20,9 @@ from vllm_omni.platforms import current_omni_platform
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate text-to-video with a LingBot-Video checkpoint.")
parser.add_argument("--model", default="robbyant/lingbot-video-dense-1.3b")
parser.add_argument("--model", default="/home/models/lingbot-video-dense-1.3b")
parser.add_argument("--prompt", default="a robotic arm picks up a red block")
parser.add_argument("--negative-prompt", default=None)
parser.add_argument("--negative-prompt", default="")
parser.add_argument("--output", default="lingbot_video_output.mp4")
parser.add_argument("--height", type=int, default=192)
parser.add_argument("--width", type=int, default=320)
@@ -58,12 +58,9 @@ def _to_frame_array(video: Any) -> np.ndarray:
if isinstance(video, torch.Tensor):
tensor = video.detach().cpu()
if tensor.dim() == 5:
tensor = tensor[0]
if tensor.dim() == 4 and tensor.shape[-1] not in (3, 4):
if tensor.shape[0] in (3, 4):
tensor = tensor.permute(1, 2, 3, 0)
elif tensor.shape[1] in (3, 4):
tensor = tensor.permute(0, 2, 3, 1)
tensor = tensor[0] if tensor.shape[1] not in (3, 4) else tensor[0].permute(1, 2, 3, 0)
elif tensor.dim() == 4 and tensor.shape[0] in (3, 4):
tensor = tensor.permute(1, 2, 3, 0)
if tensor.is_floating_point():
if float(tensor.min()) < 0.0:
tensor = tensor.clamp(-1, 1) * 0.5 + 0.5
@@ -90,12 +87,8 @@ def main() -> None:
parallel_config=DiffusionParallelConfig(),
)
prompt: dict[str, Any] = {
"prompt": args.prompt,
"modalities": ["video"],
"num_frames": args.num_frames,
}
if args.negative_prompt is not None:
prompt: dict[str, str] = {"prompt": args.prompt}
if args.negative_prompt:
prompt["negative_prompt"] = args.negative_prompt
sampling_params = OmniDiffusionSamplingParams(
@@ -1,74 +0,0 @@
# LingBot-Video
LingBot-Video uses one `LingBotVideoPipeline` for text-to-image (T2I),
text-to-video (T2V), and text-image-to-video (TI2V) generation. Both the dense
and MoE checkpoints use the same request format.
## Start the server
```bash
MODEL=robbyant/lingbot-video-dense-1.3b bash run_server.sh
```
The MoE checkpoint uses the same server and request scripts, but requires
substantially more GPU memory:
```bash
MODEL=robbyant/lingbot-video-moe-30b-a3b bash run_server.sh
```
## Text to image
The image endpoint selects T2I mode and always generates one frame:
```bash
bash run_curl_text_to_image.sh
```
The script sends a `320x192`, two-step smoke request and writes
`lingbot_t2i.png`.
## Text or text-image to video
Run the video script without an image to select T2V mode:
```bash
bash run_curl_text_image_to_video.sh
```
Pass a first-frame image to the same script to select TI2V mode:
```bash
INPUT_IMAGE=/path/to/input.png bash run_curl_text_image_to_video.sh
```
The client scripts omit the optional `model` request field, so they target
whichever dense or MoE checkpoint the server loaded. The video example uses the
lightweight `320x192`, 9-frame, two-step configuration.
Until the shared `/v1/videos` reference-image resizing is removed, TI2V target
dimensions must be sent through `extra_params`, for example
`{"size":"320x192"}`. Do not use the top-level `size`, `width`, or `height`
fields for TI2V because the serving layer currently applies those dimensions
to the reference image before the model receives it. T2V requests can continue
to use the top-level dimension fields.
LingBot video frame counts use the causal VAE `4n+1` grid. The pipeline rounds
any requested frame count upward to the next valid value. An explicit
`num_frames` takes precedence over `seconds`; otherwise, the server first
resolves `seconds * fps` and the pipeline applies the same alignment.
Official `resolution`/`ratio` presets can be sent through `extra_params`, for
example `{"resolution":"720p","ratio":"16:9"}`. The `2k` and `4k` entries
only define output dimensions; whether they run successfully depends on the
checkpoint, GPU memory, and memory optimizations available in the deployment.
For `/v1/images/generations`, the server resolves these aliases to their final
pixel dimensions before applying `--max-generated-image-size`. Requests above
the configured limit return HTTP 400 before engine dispatch. LingBot produces
one output per prompt; image requests with `n>1` are also rejected with HTTP
400.
LingBot TI2V accepts exactly one image reference. Image editing, video
references, audio references, batching, and Refiner execution are not supported
by this pipeline mode.
@@ -1,67 +0,0 @@
#!/bin/bash
set -euo pipefail
if [ -n "${INPUT_IMAGE:-}" ] && [ ! -f "${INPUT_IMAGE}" ]; then
echo "INPUT_IMAGE does not exist: ${INPUT_IMAGE}"
exit 1
fi
BASE_URL="${BASE_URL:-http://localhost:8099}"
POLL_INTERVAL="${POLL_INTERVAL:-2}"
if [ -n "${INPUT_IMAGE:-}" ]; then
default_output="lingbot_ti2v.mp4"
else
default_output="lingbot_t2v.mp4"
fi
OUTPUT_PATH="${OUTPUT_PATH:-${default_output}}"
curl_args=(
-sS -X POST "${BASE_URL}/v1/videos"
-H "Accept: application/json"
-F "prompt=the subject turns toward the camera with smooth natural motion"
-F "num_frames=9"
-F "fps=24"
-F "num_inference_steps=2"
-F "guidance_scale=3.0"
-F "flow_shift=3.0"
-F "seed=42"
)
if [ -n "${INPUT_IMAGE:-}" ]; then
curl_args+=(
-F "input_reference=@${INPUT_IMAGE}"
-F 'extra_params={"size":"320x192"}'
)
else
curl_args+=(-F "width=320" -F "height=192")
fi
create_response="$(curl "${curl_args[@]}")"
video_id="$(echo "${create_response}" | jq -er '.id')"
echo "Created video job ${video_id}"
while true; do
status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")"
status="$(echo "${status_response}" | jq -er '.status')"
case "${status}" in
queued|in_progress)
sleep "${POLL_INTERVAL}"
;;
completed)
break
;;
failed)
echo "${status_response}" | jq .
exit 1
;;
*)
echo "Unexpected video job status: ${status}"
exit 1
;;
esac
done
curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}"
echo "Saved video to ${OUTPUT_PATH}"
@@ -1,20 +0,0 @@
#!/bin/bash
set -euo pipefail
BASE_URL="${BASE_URL:-http://localhost:8099}"
OUTPUT_PATH="${OUTPUT_PATH:-lingbot_t2i.png}"
curl -sS -X POST "${BASE_URL}/v1/images/generations" \
-H "Content-Type: application/json" \
-d "{
\"prompt\": \"a red fox standing in fresh snow\",
\"size\": \"320x192\",
\"num_inference_steps\": 2,
\"guidance_scale\": 3.0,
\"flow_shift\": 3.0,
\"seed\": 42,
\"response_format\": \"b64_json\"
}" | jq -er '.data[0].b64_json' | base64 -d > "${OUTPUT_PATH}"
echo "Saved image to ${OUTPUT_PATH}"
@@ -1,10 +0,0 @@
#!/bin/bash
set -euo pipefail
MODEL="${MODEL:-robbyant/lingbot-video-dense-1.3b}"
PORT="${PORT:-8099}"
vllm serve "${MODEL}" --omni \
--model-class-name LingBotVideoPipeline \
--port "${PORT}"
+3 -3
View File
@@ -23,10 +23,10 @@ including `librosa`.
The deploy config auto-loads via `--omni`.
The default `vllm_omni/deploy/minicpmo_4_5.yaml` keeps all three stages on
logical device 0 with memory budgets of 55%, 15%, and 18%, leaving headroom
logical device 0 with memory budgets of 55%, 15%, and 15%, leaving headroom
for runtime kernels such as the HiFi-GAN vocoder's cuDNN workspace. The
profile admits at most four sequences per stage and bounds Talker context
to 4096 tokens. For throughput,
profile admits at most four sequences per stage and bounds startup video
profiling to 32 frames per video. For throughput,
`minicpmo_4_5_2gpu.yaml` gives the Thinker
GPU 0 (90%) and colocates the Talker (55%) and Code2Wav (35%) on GPU 1. That
profile admits at most four concurrent sequences per stage.
@@ -196,7 +196,6 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
args.model,
ref_audio=_ref_audio_data_url(args.ref_audio),
session_id=args.session_id,
temperature=args.temperature,
timeout_s=args.timeout_s,
)
stream_event_cursor = len(client.events.events)
@@ -384,12 +383,6 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--output-dir", default="/tmp/minicpmo_realtime_duplex_demo")
parser.add_argument("--chunk-ms", type=int, default=200)
parser.add_argument("--timeout-s", type=float, default=60.0)
parser.add_argument(
"--temperature",
type=float,
default=None,
help="Stage0 sampling temperature; response-required soft-interrupt uses 0.0.",
)
parser.add_argument("--no-realtime-pacing", action="store_true")
parser.add_argument("--require-audio", action="store_true")
return parser.parse_args()
@@ -0,0 +1,247 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Verify a ModelOpt FP8 diffusers checkpoint exported by
quantize_wan2_2_modelopt_fp8.py / quantize_hunyuanvideo_15_modelopt_fp8.py.
Three checks:
A. transformer/config.json has a sane quantization_config block.
B. transformer/*.safetensors contains FP8 (float8_e4m3fn) quantized tensors.
C. transformer disk size is materially smaller than a BF16 baseline.
Example:
python examples/quantization/check_modelopt_fp8_export.py \\
--output ./hv15-480p-modelopt-fp8
# Optional: compare disk size against a local or HF BF16 baseline.
python examples/quantization/check_modelopt_fp8_export.py \\
--output ./hv15-480p-modelopt-fp8 \\
--baseline hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
"""
from __future__ import annotations
import argparse
import json
import sys
from collections import Counter
from pathlib import Path
SUPPORTED_ALGOS = ("FP8",)
def _check_config(transformer_dir: Path) -> tuple[int, str | None]:
"""Returns (status, quant_algo). status: 0 pass, 1 fail, 2 warn."""
cfg_path = transformer_dir / "config.json"
if not cfg_path.exists():
print(f"[FAIL] {cfg_path} missing.")
return 1, None
with cfg_path.open(encoding="utf-8") as f:
cfg = json.load(f)
qc = cfg.get("quantization_config")
if not isinstance(qc, dict):
print(f"[FAIL] No `quantization_config` block in {cfg_path}.")
return 1, None
print(f"[A] quantization_config from {cfg_path}:")
print(json.dumps(qc, indent=2))
quant_algo = qc.get("quant_algo")
issues = []
if qc.get("quant_method") != "modelopt":
issues.append(f"quant_method={qc.get('quant_method')!r} (expected 'modelopt')")
if quant_algo not in SUPPORTED_ALGOS:
issues.append(
f"quant_algo={quant_algo!r} (expected one of {SUPPORTED_ALGOS} — vllm-omni adapter may not auto-detect)"
)
if issues:
print("[A] WARN — config looks incomplete:")
for issue in issues:
print(f" - {issue}")
return 2, quant_algo
print(f"[A] PASS — config looks correct (quant_algo={quant_algo}).")
return 0, quant_algo
def _read_safetensors_header(path: Path) -> dict:
"""Read the JSON header of a safetensors file. Bypass-safe — doesn't materialize tensors.
Returns {tensor_name: {'dtype': 'F8_E4M3', 'shape': [...], 'data_offsets': [...]}}.
Header dtype strings: F8_E4M3, F8_E5M2, BF16, F16, F32, F64, I8, I16, I32, I64, BOOL, U8, ...
"""
import struct
with open(path, "rb") as f:
header_len = struct.unpack("<Q", f.read(8))[0]
header = json.loads(f.read(header_len))
header.pop("__metadata__", None)
return header
def _classify_weight_scale_granularity(weight_scale_shapes: list[list[int]]) -> str:
"""Infer per-tensor vs per-channel vs per-block from sample weight_scale shapes.
ModelOpt block-wise produces shapes like `[16, 1, 16, 1]` (broadcasting dims of 1
interleaved with block-count dims). We count "meaningful" dims — ones with size > 1 —
and classify: 0 meaningful dims = per-tensor (scalar), 1 = per-channel, 2+ = per-block.
"""
if not weight_scale_shapes:
return "no weight_scale tensors found"
def meaningful_dims(shape: list[int]) -> int:
return sum(1 for d in shape if d > 1)
per_tensor = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 0)
per_channel = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 1)
per_block = sum(1 for s in weight_scale_shapes if meaningful_dims(s) >= 2)
total = len(weight_scale_shapes)
if per_tensor == total:
return "per-tensor (all scalar scales)"
if per_channel == total:
return "per-channel (1 meaningful dim)"
if per_block == total:
return "per-block (2+ meaningful dims — e.g. [M//bm, 1, N//bn, 1] for tiles)"
return f"mixed: per-tensor={per_tensor}, per-channel={per_channel}, per-block={per_block} of {total}"
def _check_safetensors(transformer_dir: Path) -> int:
"""Returns 0 on pass, 1 on fail. Reads on-disk dtype from the safetensors header."""
files = sorted(transformer_dir.glob("*.safetensors"))
if not files:
print(f"[FAIL] No *.safetensors in {transformer_dir}.")
return 1
header_dtype_counts: Counter[str] = Counter()
sample_quant_weight_keys: list[str] = []
sample_scale_keys: list[str] = []
weight_scale_shapes: list[list[int]] = []
sample_weight_scale_entries: list[tuple[str, list[int]]] = []
for f in files:
try:
header = _read_safetensors_header(f)
except Exception as exc:
print(f"[B] WARN — could not parse header of {f}: {exc}")
continue
for k, info in header.items():
dtype = info.get("dtype", "?")
header_dtype_counts[dtype] += 1
# FP8 stores weights as F8_E4M3 directly.
if dtype.startswith("F8") and k.endswith(".weight") and len(sample_quant_weight_keys) < 5:
sample_quant_weight_keys.append(k)
if k.endswith(("_scale", ".weight_scale", ".input_scale", "_scale_inv")) and len(sample_scale_keys) < 5:
sample_scale_keys.append(k)
if k.endswith(".weight_scale"):
weight_scale_shapes.append(info.get("shape", []))
if len(sample_weight_scale_entries) < 5:
sample_weight_scale_entries.append((k, info.get("shape", [])))
print(f"\n[B] On-disk dtype counts across {len(files)} safetensors file(s) (from header, not get_tensor):")
for dtype, count in sorted(header_dtype_counts.items(), key=lambda kv: -kv[1]):
marker = " <-- FP8" if dtype.startswith("F8") else ""
print(f" {dtype:10s} {count:>6d}{marker}")
quant_count = sum(c for d, c in header_dtype_counts.items() if d.startswith("F8"))
if quant_count == 0:
print("[B] FAIL — no FP8 tensors on disk. Calibration likely did not actually quantize the weights.")
return 1
print(f"[B] PASS — {quant_count} FP8 tensors stored on disk.")
if sample_quant_weight_keys:
print(f" sample quantized weight tensors: {sample_quant_weight_keys[:3]}")
if sample_scale_keys:
print(f" sample scale tensors: {sample_scale_keys[:3]}")
print(" (Note: torch's get_tensor() may return these as bf16 views on some versions —")
print(" irrelevant; vLLM's loader uses native FP8 ops.)")
# Weight-scale granularity — per-tensor (scalar) vs per-channel (1-D) vs per-block (N-D).
print(f"\n weight_scale granularity: {_classify_weight_scale_granularity(weight_scale_shapes)}")
for key, shape in sample_weight_scale_entries[:3]:
print(f" {key}: shape {shape}")
return 0
def _disk_size_gib(p: Path) -> float:
return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) / (1024**3)
def _check_size_vs_baseline(transformer_dir: Path, baseline: str | None) -> int:
"""Returns 0 always (informational only)."""
quant_size = _disk_size_gib(transformer_dir)
print(f"\n[C] FP8 transformer disk size: {quant_size:.2f} GiB")
if baseline is None:
print("[C] SKIP — pass --baseline <path or HF id> to compare against BF16.")
return 0
baseline_path = Path(baseline)
if not baseline_path.exists():
# Try HF download.
try:
from huggingface_hub import snapshot_download
except ImportError:
print("[C] SKIP — huggingface_hub not installed and baseline not a local path.")
return 0
print(f" Downloading baseline transformer from HF: {baseline}")
baseline_path = Path(snapshot_download(baseline, allow_patterns=["transformer/*"]))
bf16_dir = baseline_path / "transformer" if (baseline_path / "transformer").exists() else baseline_path
bf16_size = _disk_size_gib(bf16_dir)
if bf16_size == 0:
print(f"[C] WARN — baseline transformer dir empty: {bf16_dir}")
return 0
# Expected reduction for FP8: ~50%.
min_reduction = 30
reduction = (1 - quant_size / bf16_size) * 100
print(f"[C] BF16 baseline transformer disk size: {bf16_size:.2f} GiB ({bf16_dir})")
print(f"[C] Disk reduction: {reduction:.1f}% (FP8 transformer is {quant_size / bf16_size:.0%} of BF16)")
if reduction < min_reduction:
print(
f"[C] WARN — FP8 should typically reduce disk by ~40-50%; <{min_reduction}% suggests partial quantization."
)
return 0
def main() -> None:
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--output", required=True, help="Path to the exported ModelOpt FP8 checkpoint root.")
p.add_argument(
"--baseline",
default=None,
help="Optional BF16 baseline (local diffusers dir or HF id) for disk-size comparison.",
)
args = p.parse_args()
out_root = Path(args.output).expanduser().resolve()
transformer_dir = out_root / "transformer"
if not transformer_dir.exists():
print(f"[FAIL] {transformer_dir} does not exist.")
sys.exit(1)
print(f"Checking: {out_root}\n")
fail = 0
config_status, _quant_algo = _check_config(transformer_dir)
fail |= config_status
fail |= _check_safetensors(transformer_dir)
_check_size_vs_baseline(transformer_dir, args.baseline)
print()
if fail == 0:
print("=" * 60)
print("ALL CHECKS PASSED — checkpoint looks ready for vllm-omni serving.")
elif fail == 1:
print("=" * 60)
print("FAILURES detected — calibration may need to be re-run.")
sys.exit(1)
else:
print("=" * 60)
print("WARNINGS only — checkpoint may serve but with caveats. See [A] above.")
if __name__ == "__main__":
main()
@@ -0,0 +1,448 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Quantize HunyuanVideo-1.5 (480p T2V) to a ModelOpt FP8 Hugging Face checkpoint.
Calibrates the DiT transformer using a small video prompt set and exports a
diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
Layers kept full precision match the #2728 / #2795 pattern: modulation,
AdaLayerNorm, entry/exit projections, embeddings, the token refiner path,
and final proj_out. MHA quantizers are off by default; HV-1.5 self-attention
empirically degrades under FP8 (see #2920 ablation).
Example:
python examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py \\
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\
--output ./hv15-480p-modelopt-fp8 \\
--overwrite
"""
from __future__ import annotations
import argparse
import copy
import json
import re
import shutil
import sys
from pathlib import Path
from typing import Any
import torch
from diffusers import DiffusionPipeline
DEFAULT_PROMPTS = [
"A dog running across a field of golden wheat.",
"An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
"A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
"A crackling campfire at night under a starry sky, sparks rising into the dark.",
"An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
"A close-up of a blooming rose covered in morning dew, soft natural light.",
"A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
"A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
]
def _build_parser() -> argparse.ArgumentParser:
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--model", required=True, help="Input HV-1.5 diffusers directory or HF id.")
p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
p.add_argument("--height", type=int, default=480)
p.add_argument("--width", type=int, default=832)
p.add_argument(
"--num-frames",
type=int,
default=33,
help="Frames per calibration sample. 33 matches the typical short benchmark.",
)
p.add_argument("--guidance-scale", type=float, default=6.0)
p.add_argument(
"--calib-steps",
type=int,
default=10,
help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
)
p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
p.add_argument("--seed", type=int, default=42)
p.add_argument(
"--prompt",
action="append",
default=[],
help="Custom calibration prompt. Repeat to provide multiple.",
)
p.add_argument(
"--quantize-mha",
action="store_true",
help="Enable FP8 attention K/V/softmax quantizers. Off by default — empirically degrades HV-1.5 video output.",
)
p.add_argument(
"--weight-block-size",
type=str,
default=None,
help="Per-block weight quantization as 'M,N' (e.g. '128,128' for 128x128 tiles). "
"Default: per-tensor (one scale per linear). Block-wise typically gives tighter quality at "
"negligible memory cost. Static FP8 is exempt from upstream vLLM's online block-wise gate.",
)
p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
return p
def _parse_block_size(spec: str | None) -> list[int] | None:
if spec is None:
return None
parts = [int(x) for x in spec.split(",") if x.strip()]
if len(parts) != 2:
raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
return parts
def _require_modelopt() -> Any:
try:
import modelopt.torch.quantization as mtq
except ModuleNotFoundError as exc:
raise SystemExit(
"NVIDIA ModelOpt is not installed. Install with:\n"
" pip install 'nvidia-modelopt[all]'\n"
f"Original error: {exc}"
) from exc
return mtq
def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
model_path = args.model
output_dir = Path(args.output).expanduser().resolve()
if output_dir.exists():
if not args.overwrite:
raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
shutil.rmtree(output_dir)
return model_path, output_dir
def _select_dtype(name: str) -> torch.dtype:
return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
def _build_prompts(args: argparse.Namespace) -> list[str]:
prompts = args.prompt or DEFAULT_PROMPTS
if args.calib_size <= 0:
raise SystemExit("--calib-size must be positive.")
if len(prompts) < args.calib_size:
repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
prompts = (prompts * repeats)[: args.calib_size]
return prompts[: args.calib_size]
# Layers to KEEP at full precision (mirror of the #2920 wiring + #2728/#2795 skip pattern).
# - x_embedder, image_embedder, context_embedder*, time_embed*, cond_type_embed: entry/embedding
# - norm_out, norm1*.linear, norm1_context*.linear, norm2*, norm2_context*: AdaLayerNorm modulation
# - proj_out: final output projection
# - token_refiner*: text-encoder refinement uses diffusers raw nn.Linear
def _filter_func_hv15(name: str) -> bool:
pattern = re.compile(
r"(proj_out.*|"
r".*(x_embedder|image_embedder|context_embedder|context_embedder_2|"
r"time_embed|cond_type_embed|"
r"norm_out|norm1\.linear|norm1_context\.linear|norm2|norm2_context|"
r"token_refiner).*)"
)
return pattern.match(name) is not None
def _mha_filter_func(name: str) -> bool:
pattern = re.compile(
r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
)
return pattern.match(name) is not None
def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
if not hasattr(mtq, "disable_quantizer"):
return
mtq.disable_quantizer(backbone, _filter_func_hv15)
if not quantize_mha:
mtq.disable_quantizer(backbone, _mha_filter_func)
def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
if hasattr(pipe, "set_progress_bar_config"):
pipe.set_progress_bar_config(disable=True)
pipe.to("cuda")
return pipe
def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
generator = torch.Generator(device="cuda")
# Try to set guidance on the pipeline's guider object up front (modern
# diffusers HV-1.5 uses a Guider abstraction, not a per-call kwarg). Falls
# back silently — calibration uses whatever default the pipeline ships with.
guider = getattr(pipe, "guider", None)
if guider is not None and hasattr(guider, "guidance_scale"):
try:
guider.guidance_scale = args.guidance_scale
except Exception:
pass
base_kwargs = dict(
height=args.height,
width=args.width,
num_frames=args.num_frames,
num_inference_steps=args.calib_steps,
output_type="latent",
)
def forward_loop(*_unused_args, **_unused_kwargs) -> None:
with torch.inference_mode():
for idx, prompt in enumerate(prompts):
generator.manual_seed(args.seed + idx)
# Try with guidance_scale first; fall back without on TypeError
# for pipelines (like HV-1.5) that take CFG via guider config.
try:
pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
except TypeError as exc:
if "guidance_scale" not in str(exc):
raise
pipe(prompt=prompt, generator=generator, **base_kwargs)
return forward_loop
def _summarize_export(output_dir: Path) -> None:
cfg_path = output_dir / "transformer" / "config.json"
if not cfg_path.exists():
print(f"[warn] {cfg_path} missing.", file=sys.stderr)
return
with cfg_path.open(encoding="utf-8") as f:
cfg = json.load(f)
qc = cfg.get("quantization_config")
if not isinstance(qc, dict):
print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
return
print("Export summary:")
print(f" quant_method: {qc.get('quant_method')}")
print(f" quant_algo: {qc.get('quant_algo')}")
producer = qc.get("producer")
if isinstance(producer, dict):
print(f" producer: {producer.get('name')} {producer.get('version')}")
print(f" config path: {cfg_path}")
def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
"""Convert in-memory weights of quantized modules to actual FP8 storage.
`export_hf_checkpoint` skips this step for unknown model types (HV-1.5 isn't
in ModelOpt's recognized-model registry), so we must call the per-weight
export helper ourselves. Same workaround as the HunyuanImage-3 calibration
helper.
"""
from modelopt.torch.export.quant_utils import (
QUANTIZATION_NONE,
get_quantization_format,
quantizer_attr_names,
weight_attr_names,
)
from modelopt.torch.export.unified_export_hf import _export_quantized_weight
exported = 0
for name, module in backbone.named_modules():
try:
quantization_format = get_quantization_format(module)
except Exception as exc:
print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
continue
if quantization_format == QUANTIZATION_NONE:
continue
for weight_name in weight_attr_names(module):
quantizer_attrs = quantizer_attr_names(weight_name)
weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
continue
_export_quantized_weight(module, dtype, weight_name)
exported += 1
return exported
def _hv15_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
"""Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913).
Same shape as the HunyuanImage-3 author's _hunyuan_quant_config(). When
`weight_block_size` is given, advertise block-wise weight quantization in
the saved metadata (so consumers know to expect multi-element scale tensors).
"""
weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
if weight_block_size is not None:
weights_cfg["strategy"] = "block"
weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
return {
"config_groups": {
"group_0": {
"input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
"weights": weights_cfg,
"targets": ["Linear"],
}
},
"ignore": [
"context_embedder*",
"context_embedder_2*",
"cond_type_embed*",
"image_embedder*",
"norm1.linear*",
"norm1_context.linear*",
"norm2*",
"norm2_context*",
"norm_out*",
"proj_out*",
"time_embed*",
"token_refiner*",
"x_embedder*",
],
"producer": {"name": "modelopt"},
"quant_algo": "FP8",
"quant_method": "modelopt",
}
def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
"""Inject quant_algo: FP8 + config_groups into transformer/config.json so
vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8."""
cfg_path = output_dir / "transformer" / "config.json"
with cfg_path.open(encoding="utf-8") as f:
cfg = json.load(f)
new_qc = _hv15_quant_config_block(weight_block_size=weight_block_size)
existing = cfg.get("quantization_config")
if isinstance(existing, dict):
producer = existing.get("producer")
if isinstance(producer, dict):
new_qc["producer"] = producer
cfg["quantization_config"] = new_qc
with cfg_path.open("w", encoding="utf-8") as f:
json.dump(cfg, f, indent=2)
def _save_pipeline_with_fp8_transformer(
pipe: DiffusionPipeline,
model_path: str,
output_dir: Path,
max_shard_size: str = "5GB",
) -> None:
"""Save the pipeline with the (now FP8) transformer.
Copies the source directory verbatim except for `transformer/`, then
saves the transformer with quantizers hidden so the state dict contains
only the FP8 weights + scale tensors.
"""
from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
src = Path(model_path)
if not src.exists():
from huggingface_hub import snapshot_download
src = Path(snapshot_download(model_path))
if output_dir.exists():
shutil.rmtree(output_dir)
shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer"))
transformer_out = output_dir / "transformer"
# `hide_quantizers_from_state_dict` walks named_modules(); pass the actual
# nn.Module (transformer), not the diffusers Pipeline wrapper.
with hide_quantizers_from_state_dict(pipe.transformer):
pipe.transformer.save_pretrained(
str(transformer_out),
safe_serialization=True,
max_shard_size=max_shard_size,
)
def main() -> None:
args = _build_parser().parse_args()
if not torch.cuda.is_available():
raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
mtq = _require_modelopt()
model_path, output_dir = _ensure_paths(args)
dtype = _select_dtype(args.dtype)
prompts = _build_prompts(args)
print("Quantization plan:")
weight_block_size = _parse_block_size(args.weight_block_size)
print(f" input: {args.model}")
print(f" output: {output_dir}")
print(f" dtype: {dtype}")
print(f" height/width: {args.height}x{args.width}")
print(f" num_frames: {args.num_frames}")
print(f" calib_size: {len(prompts)}")
print(f" calib_steps: {args.calib_steps}")
print(f" quantize_mha: {args.quantize_mha}")
print(
f" weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
)
pipe = _load_pipeline(model_path, dtype)
backbone = pipe.transformer
quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
if weight_block_size is not None:
# Switch from per-tensor (default) to block-wise weight quantization.
# ModelOpt's wildcard "*weight_quantizer" matches every linear's weight quantizer.
quant_config["quant_cfg"]["*weight_quantizer"] = {
"num_bits": (4, 3), # E4M3 (FP8 weights, same as default)
"block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
}
print(
f" -> overriding weight quantizer with block_sizes={weight_block_size} "
f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
)
forward_loop = _build_forward_loop(pipe, args, prompts)
quantized = mtq.quantize(backbone, quant_config, forward_loop)
if quantized is not None:
pipe.transformer = quantized
backbone = quantized
_disable_known_problematic_quantizers(mtq, backbone, quantize_mha=args.quantize_mha)
print("\nForcing FP8 weight serialization (HV-1.5 isn't in ModelOpt's recognized-model registry,")
print("so we have to call the per-weight export helper ourselves)...")
exported = _force_export_quantized_weights(backbone, dtype)
print(f" -> {exported} weights converted to FP8 in memory")
if exported == 0:
raise SystemExit(
"No quantized weights were exported. Calibration may have skipped every layer "
"(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
"weight quantizers."
)
print("\nSaving pipeline with FP8 transformer...")
_save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
_patch_quant_config(output_dir, weight_block_size=weight_block_size)
print(f"Saved to: {output_dir}")
_summarize_export(output_dir)
print("\nNext: validate the checkpoint with vllm-omni:")
print(
" python examples/offline_inference/text_to_video/text_to_video.py \\\n"
f" --model {output_dir} \\\n"
" --quantization fp8 \\\n"
" --prompt 'A dog running across a field of golden wheat.' \\\n"
f" --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
" --num-inference-steps 30 --guidance-scale 6.0 --seed 42 \\\n"
" --output outputs/hv15_modelopt_fp8.mp4 \\\n"
" --enforce-eager"
)
print(
"\n (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
"checkpoint's config.json has modelopt metadata.)"
)
if __name__ == "__main__":
main()
@@ -0,0 +1,471 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Quantize Wan2.2 (TI2V-5B, 704x1280 T2V) to a ModelOpt FP8 Hugging Face checkpoint.
Calibrates the DiT transformer using a small video prompt set and exports a
diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
Layers kept full precision match the #2728 / #2795 pattern: condition embedder
(time/text/image), patch embedding, modulation (scale_shift_table), final
norm + proj_out, and sequence-parallel helpers. All attention + FFN linears
are quantized — static calibration handles the numerics that online FP8
couldn't (see #2920 ablation).
Default target is `Wan-AI/Wan2.2-TI2V-5B-Diffusers`, the dense 5B variant that
fits 80GB BF16. The A14B MoE variants need 2+ GPUs and are out of scope here.
Example:
python examples/quantization/quantize_wan2_2_modelopt_fp8.py \\
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \\
--output ./wan22-ti2v-modelopt-fp8 \\
--overwrite
"""
from __future__ import annotations
import argparse
import copy
import json
import re
import shutil
import sys
from pathlib import Path
from typing import Any
import torch
from diffusers import DiffusionPipeline
DEFAULT_PROMPTS = [
"A dog running across a field of golden wheat.",
"An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
"A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
"A crackling campfire at night under a starry sky, sparks rising into the dark.",
"An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
"A close-up of a blooming rose covered in morning dew, soft natural light.",
"A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
"A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
]
def _build_parser() -> argparse.ArgumentParser:
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--model", required=True, help="Input Wan2.2 diffusers directory or HF id.")
p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
p.add_argument("--height", type=int, default=704, help="Calibration video height (Wan2.2 TI2V-5B native: 704).")
p.add_argument("--width", type=int, default=1280, help="Calibration video width (Wan2.2 TI2V-5B native: 1280).")
p.add_argument(
"--num-frames",
type=int,
default=49,
help="Frames per calibration sample. 49 matches the typical short benchmark; "
"use 17 to reduce memory pressure during calibration.",
)
p.add_argument("--guidance-scale", type=float, default=5.0)
p.add_argument(
"--calib-steps",
type=int,
default=10,
help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
)
p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
p.add_argument("--seed", type=int, default=42)
p.add_argument(
"--prompt",
action="append",
default=[],
help="Custom calibration prompt. Repeat to provide multiple.",
)
p.add_argument(
"--quantize-mha",
action="store_true",
help="Enable FP8 attention K/V/softmax quantizers. Off by default — Wan2.2's long attention "
"sequences amplified FP8 drift in the online ablation (see #2920).",
)
p.add_argument(
"--weight-block-size",
type=str,
default=None,
help="Per-block weight quantization as 'M,N' (e.g. '128,128'). Default per-tensor. "
"Note: vllm-omni's ModelOpt adapter may not yet dispatch block-wise scales — check #2924 "
"for the HV-1.5 investigation status before relying on this.",
)
p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
return p
def _parse_block_size(spec: str | None) -> list[int] | None:
if spec is None:
return None
parts = [int(x) for x in spec.split(",") if x.strip()]
if len(parts) != 2:
raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
return parts
def _require_modelopt() -> Any:
try:
import modelopt.torch.quantization as mtq
except ModuleNotFoundError as exc:
raise SystemExit(
"NVIDIA ModelOpt is not installed. Install with:\n"
" pip install 'nvidia-modelopt[all]'\n"
f"Original error: {exc}"
) from exc
return mtq
def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
model_path = args.model
output_dir = Path(args.output).expanduser().resolve()
if output_dir.exists():
if not args.overwrite:
raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
shutil.rmtree(output_dir)
return model_path, output_dir
def _select_dtype(name: str) -> torch.dtype:
return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
def _build_prompts(args: argparse.Namespace) -> list[str]:
prompts = args.prompt or DEFAULT_PROMPTS
if args.calib_size <= 0:
raise SystemExit("--calib-size must be positive.")
if len(prompts) < args.calib_size:
repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
prompts = (prompts * repeats)[: args.calib_size]
return prompts[: args.calib_size]
# Layers to KEEP at full precision. Wan2.2's module naming:
# - condition_embedder: time_embedder, time_proj, text_embedder, image_embedder (I2V)
# - patch_embedding: Conv3dLayer (already not Linear, belt-and-suspenders skip)
# - scale_shift_table: nn.Parameter modulation (not Linear, but pattern guard)
# - norm_out: AdaLayerNorm final
# - proj_out: final nn.Linear
# - timestep_proj_prepare / output_scale_shift_prepare: SP helpers
def _filter_func_wan22(name: str) -> bool:
pattern = re.compile(
r"(proj_out.*|"
r".*(condition_embedder|patch_embedding|"
r"norm_out|scale_shift_table|"
r"timestep_proj_prepare|output_scale_shift_prepare).*)"
)
return pattern.match(name) is not None
def _mha_filter_func(name: str) -> bool:
pattern = re.compile(
r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
)
return pattern.match(name) is not None
def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
if not hasattr(mtq, "disable_quantizer"):
return
mtq.disable_quantizer(backbone, _filter_func_wan22)
if not quantize_mha:
mtq.disable_quantizer(backbone, _mha_filter_func)
def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
if hasattr(pipe, "set_progress_bar_config"):
pipe.set_progress_bar_config(disable=True)
pipe.to("cuda")
return pipe
def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
generator = torch.Generator(device="cuda")
# Try setting guidance on the pipeline's guider if present (newer diffusers APIs).
guider = getattr(pipe, "guider", None)
if guider is not None and hasattr(guider, "guidance_scale"):
try:
guider.guidance_scale = args.guidance_scale
except Exception:
pass
base_kwargs = dict(
height=args.height,
width=args.width,
num_frames=args.num_frames,
num_inference_steps=args.calib_steps,
output_type="latent",
)
def forward_loop(*_unused_args, **_unused_kwargs) -> None:
with torch.inference_mode():
for idx, prompt in enumerate(prompts):
generator.manual_seed(args.seed + idx)
# Try with guidance_scale first; fall back without on TypeError
# for pipelines that take CFG via guider config only.
try:
pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
except TypeError as exc:
if "guidance_scale" not in str(exc):
raise
pipe(prompt=prompt, generator=generator, **base_kwargs)
return forward_loop
def _summarize_export(output_dir: Path) -> None:
cfg_path = output_dir / "transformer" / "config.json"
if not cfg_path.exists():
print(f"[warn] {cfg_path} missing.", file=sys.stderr)
return
with cfg_path.open(encoding="utf-8") as f:
cfg = json.load(f)
qc = cfg.get("quantization_config")
if not isinstance(qc, dict):
print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
return
print("Export summary:")
print(f" quant_method: {qc.get('quant_method')}")
print(f" quant_algo: {qc.get('quant_algo')}")
producer = qc.get("producer")
if isinstance(producer, dict):
print(f" producer: {producer.get('name')} {producer.get('version')}")
print(f" config path: {cfg_path}")
def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
"""Convert in-memory weights of quantized modules to actual FP8 storage.
`export_hf_checkpoint` skips this step for unknown model types (Wan2.2 isn't
in ModelOpt's recognized-model registry), so we must call the per-weight
export helper ourselves. Same workaround as the HunyuanVideo-1.5 / HunyuanImage-3
calibration helpers.
"""
from modelopt.torch.export.quant_utils import (
QUANTIZATION_NONE,
get_quantization_format,
quantizer_attr_names,
weight_attr_names,
)
from modelopt.torch.export.unified_export_hf import _export_quantized_weight
exported = 0
for name, module in backbone.named_modules():
try:
quantization_format = get_quantization_format(module)
except Exception as exc:
print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
continue
if quantization_format == QUANTIZATION_NONE:
continue
for weight_name in weight_attr_names(module):
quantizer_attrs = quantizer_attr_names(weight_name)
weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
continue
_export_quantized_weight(module, dtype, weight_name)
exported += 1
return exported
def _wan22_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
"""Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913)."""
weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
if weight_block_size is not None:
weights_cfg["strategy"] = "block"
weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
return {
"config_groups": {
"group_0": {
"input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
"weights": weights_cfg,
"targets": ["Linear"],
}
},
"ignore": [
"condition_embedder*",
"norm_out*",
"output_scale_shift_prepare*",
"patch_embedding*",
"proj_out*",
"scale_shift_table*",
"timestep_proj_prepare*",
],
"producer": {"name": "modelopt"},
"quant_algo": "FP8",
"quant_method": "modelopt",
}
def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
"""Inject quant_algo: FP8 + config_groups into each transformer's config.json
so vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8.
Patches both ``transformer`` and ``transformer_2`` (A14B) when present."""
for sub in ("transformer", "transformer_2"):
cfg_path = output_dir / sub / "config.json"
if not cfg_path.exists():
continue
with cfg_path.open(encoding="utf-8") as f:
cfg = json.load(f)
new_qc = _wan22_quant_config_block(weight_block_size=weight_block_size)
existing = cfg.get("quantization_config")
if isinstance(existing, dict):
producer = existing.get("producer")
if isinstance(producer, dict):
new_qc["producer"] = producer
cfg["quantization_config"] = new_qc
with cfg_path.open("w", encoding="utf-8") as f:
json.dump(cfg, f, indent=2)
def _list_transformers(pipe: DiffusionPipeline) -> list[tuple[str, torch.nn.Module]]:
"""Return [(attr_name, module)] for every DiT on the pipeline.
Wan2.2 A14B has both ``transformer`` (low-noise) and ``transformer_2``
(high-noise). TI2V-5B has only ``transformer``.
"""
out = []
for attr in ("transformer", "transformer_2"):
mod = getattr(pipe, attr, None)
if mod is not None:
out.append((attr, mod))
return out
def _save_pipeline_with_fp8_transformer(
pipe: DiffusionPipeline,
model_path: str,
output_dir: Path,
max_shard_size: str = "5GB",
) -> None:
"""Copy source dir verbatim minus transformer dirs, then save each quantized DiT.
Handles A14B's dual transformer (``transformer`` + ``transformer_2``) as well
as the single-DiT TI2V-5B.
"""
from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
src = Path(model_path)
if not src.exists():
from huggingface_hub import snapshot_download
src = Path(snapshot_download(model_path))
if output_dir.exists():
shutil.rmtree(output_dir)
shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer", "transformer_2"))
for attr, backbone in _list_transformers(pipe):
transformer_out = output_dir / attr
# Pass the nn.Module (transformer), not the Pipeline wrapper.
with hide_quantizers_from_state_dict(backbone):
backbone.save_pretrained(
str(transformer_out),
safe_serialization=True,
max_shard_size=max_shard_size,
)
def main() -> None:
args = _build_parser().parse_args()
if not torch.cuda.is_available():
raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
mtq = _require_modelopt()
model_path, output_dir = _ensure_paths(args)
dtype = _select_dtype(args.dtype)
prompts = _build_prompts(args)
weight_block_size = _parse_block_size(args.weight_block_size)
print("Quantization plan:")
print(f" input: {args.model}")
print(f" output: {output_dir}")
print(f" dtype: {dtype}")
print(f" height/width: {args.height}x{args.width}")
print(f" num_frames: {args.num_frames}")
print(f" calib_size: {len(prompts)}")
print(f" calib_steps: {args.calib_steps}")
print(f" quantize_mha: {args.quantize_mha}")
print(
f" weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
)
pipe = _load_pipeline(model_path, dtype)
transformers = _list_transformers(pipe)
if not transformers:
raise SystemExit("Pipeline has no transformer or transformer_2 attribute.")
print(f" found {len(transformers)} transformer(s): {', '.join(a for a, _ in transformers)}")
quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
if weight_block_size is not None:
quant_config["quant_cfg"]["*weight_quantizer"] = {
"num_bits": (4, 3),
"block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
}
print(
f" -> overriding weight quantizer with block_sizes={weight_block_size} "
f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
)
forward_loop = _build_forward_loop(pipe, args, prompts)
# Quantize each DiT. The forward_loop runs the full pipeline, so quantizers
# on every attached transformer (A14B: transformer + transformer_2) get
# calibrated during the same passes.
for attr, backbone in transformers:
print(f"\nQuantizing {attr}...")
quantized = mtq.quantize(backbone, quant_config, forward_loop)
if quantized is not None:
setattr(pipe, attr, quantized)
_disable_known_problematic_quantizers(mtq, getattr(pipe, attr), quantize_mha=args.quantize_mha)
# Export weights to FP8 only AFTER all transformers are calibrated. Exporting
# inside the loop would convert transformer's weights to real FP8, and the
# next transformer's forward_loop (full pipeline) would then fake-quantize an
# already-FP8 weight -> "fake_e4m3fy not implemented for Float8_e4m3fn".
print("\nForcing FP8 weight serialization (Wan2.2 isn't in ModelOpt's recognized-model registry,")
print("so we have to call the per-weight export helper ourselves)...")
total_exported = 0
for attr, backbone in _list_transformers(pipe):
exported = _force_export_quantized_weights(backbone, dtype)
print(f" -> {exported} weights converted to FP8 in {attr}")
total_exported += exported
if total_exported == 0:
raise SystemExit(
"No quantized weights were exported. Calibration may have skipped every layer "
"(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
"weight quantizers."
)
print("\nSaving pipeline with FP8 transformer(s)...")
_save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
_patch_quant_config(output_dir, weight_block_size=weight_block_size)
print(f"Saved to: {output_dir}")
_summarize_export(output_dir)
print("\nNext: validate the checkpoint with vllm-omni:")
print(
" python examples/offline_inference/text_to_video/text_to_video.py \\\n"
f" --model {output_dir} \\\n"
" --quantization fp8 \\\n"
" --prompt 'A dog running across a field of golden wheat.' \\\n"
f" --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
" --num-inference-steps 30 --guidance-scale 5.0 --seed 42 \\\n"
" --output outputs/wan22_modelopt_fp8.mp4"
)
print(
"\n (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
"checkpoint's config.json has modelopt metadata.)"
)
if __name__ == "__main__":
main()
-8
View File
@@ -89,14 +89,6 @@ local = [
"FlagEmbedding",
]
# LongCat-Video-Avatar A2V/AI2V runtime dependencies for vocal
# separation and loudness normalization.
longcat-video-avatar = [
"audio-separator==0.44.2",
"pyloudnorm==0.1.1",
"soundfile>=0.13.1",
]
# Optional Blackwell-only fused-bias FP8 GEMM; package is `quack-kernels` (imported
# as `quack`), not the unrelated `quack`. See docs/user_guide/quantization/fp8.md.
quack = [
-169
View File
@@ -1,169 +0,0 @@
# MiniMax-H3 on RTX 4090
This recipe uses BF16 weights, tiled VAE decode, tensor parallelism, optional
Ulysses sequence parallelism, and distributed layerwise offload (DLO). It is the
24 GiB sibling of [MiniMax-H3-5090.md](MiniMax-H3-5090.md): the two-GPU topology
matches that recipe, but the resident DiT-block count is lower so the per-rank
peak fits in 24 GiB. Lower resident counts reduce HBM use and increase
CPU-to-GPU transfer time.
## Capacity requirements
| Resource | Two RTX 4090s | Four RTX 4090s |
| --- | ---: | ---: |
| GPU HBM | 24 GiB per GPU | 24 GiB per GPU |
| Checkpoint storage | 135 GiB per partition | 135 GiB per partition |
| Available system RAM | 200 GiB minimum | 200 GiB minimum |
| Recommended system RAM | 384 GiB | 384 GiB |
`FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one
server at a time. DLO keeps rank-local weights in pinned host memory; increasing
`--dlo-resident-layers` improves latency but does **not** reduce host RAM in the
current implementation because resident layers retain pinned CPU master copies.
A single RTX 4090 is not covered here. The one-GPU DLO profile in the RTX 5090
recipe peaked at 26.50 GiB with 12 resident layers, which exceeds 24 GiB. On one
4090, use the model-level CPU offload command in
[MiniMax-H3.md](MiniMax-H3.md#single-gpu-accuracy-and-memory-first) instead, or
lower the resident count and re-measure before trusting it.
> **Modular H3:** after #5720 lands, preserve this recipe's one-partition
> behavior with `--task-type fl2va` or `--task-type ref2va`.
## Two RTX 4090s: 1024x576, 5 seconds
Use TP2 and 12 resident DiT layers. This is the `rtx4090` profile from the main
recipe. On two RTX 4090s it peaked at 15,620 MiB per rank for T2VA and
14,918 MiB per rank for Ref2VA, leaving roughly 9 GiB of headroom per card. See
[Target-hardware validation](#target-hardware-validation) for the full numbers.
```bash
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400
CUDA_VISIBLE_DEVICES=0,1 vllm serve /path/to/MiniMax-H3/FL2VA \
--omni --trust-remote-code --host 0.0.0.0 --port 8000 \
--num-gpus 2 --tensor-parallel-size 2 --text-encoder-tp-size 2 \
--usp 1 --ring 1 --vae-patch-parallel-size 2 \
--vae-parallel-mode tile --vae-use-tiling \
--enable-distributed-layerwise-offload --dlo-no-use-allgather \
--dlo-resident-layers 12 --enforce-eager \
--diffusion-attention-backend CUDNN_ATTN
```
For Ref2VA, stop the FL2VA server and restart the same command with
`/path/to/MiniMax-H3/Ref2VA`. Ref2VA reference image count and prompt length can
increase activation memory; begin with one request at a time.
`VLLM_OMNI_VIDEO_SYNC_TIMEOUT` matters more here than on larger cards. DLO
streams non-resident DiT blocks over PCIe on every denoising step, so a 60-step
request takes several minutes and would otherwise hit the 600-second default.
## Four RTX 4090s: 1024x576, 5 seconds
Keep TP2 and raise Ulysses sequence parallel to 2 so the world size is
`TP × USP = 4`. Text-encoder TP and VAE patch parallel follow the GPU count.
Resident layers stay at 12: adding USP does not shard DiT weights further, so
per-GPU peak HBM stays close to the two-GPU profile while wall-clock improves.
```bash
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400
CUDA_VISIBLE_DEVICES=0,1,2,3 vllm serve /path/to/MiniMax-H3/FL2VA \
--omni --trust-remote-code --host 0.0.0.0 --port 8000 \
--num-gpus 4 --tensor-parallel-size 2 --text-encoder-tp-size 4 \
--usp 2 --ring 1 --vae-patch-parallel-size 4 \
--vae-parallel-mode tile --vae-use-tiling \
--enable-distributed-layerwise-offload --dlo-no-use-allgather \
--dlo-resident-layers 12 --enforce-eager \
--diffusion-attention-backend CUDNN_ATTN
```
For Ref2VA, stop the FL2VA server and restart the same command with
`/path/to/MiniMax-H3/Ref2VA`. Pin four dedicated GPUs; if the launcher exposes
fewer than four devices, startup fails with
`Stage 0 requires 4 device(s) based on parallel_config`.
## Target-hardware validation
Measured on RTX 4090 (24,564 MiB each, driver 580.126.09) with the serve
commands above, at vLLM-Omni `0.26.1.dev55+g81b48e83e`, vLLM `0.26.0`, and
PyTorch `2.11.0+cu130`. GPUs used by each server were dedicated. `Client E2E`
and `Peak per GPU` come from the `/v1/videos/sync` response headers
`x-inference-time-s` and `x-peak-memory-mb`; the latter is the rank-0 CUDA
reserved high-water mark, so it is a per-GPU figure.
| GPUs | Topology | Task | Shape | Frames | Steps | Client E2E | Peak per GPU | Output validation |
| ---: | --- | --- | ---: | ---: | ---: | ---: | ---: | --- |
| 2 | TP2 × USP1 | T2VA (`FL2VA`) | 1024x576 | 124 at 24 FPS | 60 | 7 min 9 s | 15.3 GiB | H.264 video + 32 kHz stereo AAC; full `ffmpeg` decode passed |
| 2 | TP2 × USP1 | Ref2VA | 1024x576 | 124 at 24 FPS | 60 | 14 min 52 s | 14.6 GiB | H.264 video + 32 kHz stereo AAC; full `ffmpeg` decode passed |
| 4 | TP2 × USP2 | T2VA (`FL2VA`) | 1024x576 | 124 at 24 FPS | 60 | 4 min 29 s | 15.2 GiB | H.264 video + 32 kHz stereo AAC; full `ffmpeg` decode passed |
| 4 | TP2 × USP2 | Ref2VA | 1024x576 | 124 at 24 FPS | 60 | 9 min 5 s | 16.1 GiB | H.264 video + 32 kHz stereo AAC; successful `/v1/videos/sync` 200 |
All rows use `seed=1101`, `flow_shift=12`, and `audio_flow_shift=3.0`.
Two-GPU repeats: T2VA measured 434.9 s and 429.3 s; Ref2VA measured 891.9 s and
892.1 s. Peak memory was identical across repeats for both tasks, and the two
outputs of each pair agreed to within 115 bytes.
Four-GPU T2VA repeats: 273.8 s and 269.5 s, peaking at 15,560 MiB and
15,612 MiB. Four-GPU Ref2VA is a single successful end-to-end run (545.2 s,
16,448 MiB). Follow-up Ref2VA requests on this build completed diffusion but
then hit the engine's hardcoded 30 s async output wait
(`_ASYNC_OUTPUT_TIMEOUT` in `diffusion_engine.py`) during post-compute D2H,
so they are not reported as repeats.
Relative to the two-GPU baseline on the same shape, four-GPU TP2×USP2 was about
1.6× faster for both tasks, while per-GPU peak HBM stayed in the mid-15 GiB
range because USP does not further shard resident DiT weights. These are
single-request validation runs, not concurrent throughput benchmarks.
## Request examples
`t2va` requires an explicit `aspect_ratio`. Without it the request fails with
`t2va requires an explicit aspect_ratio`, even when `width` and `height` are set.
Text to video and audio, against the `FL2VA` server:
```bash
curl -sS -D headers_t2va.txt -o out_t2va.mp4 --max-time 14400 \
-X POST "http://127.0.0.1:8000/v1/videos/sync" \
-F 'prompt=傍晚的小厨房,夕阳余晖从窗边洒进来,旧木桌上放着洗到一半的马克杯和起雾的玻璃瓶,悬挂的抹布轻轻晃动。画面带有手持拍摄的轻微晃动和逆光曝光波动,生活感十足。环境音是安静的厨房底噪。' \
-F 'aspect_ratio=16:9' \
-F 'width=1024' -F 'height=576' -F 'fps=24' \
-F 'num_inference_steps=60' -F 'flow_shift=12' -F 'seed=1101' \
-F 'extra_params={"task":"t2va","duration":5,"audio_flow_shift":3.0}'
```
Reference to video and audio, against the `Ref2VA` server:
```bash
curl -sS -D headers_ref2va.txt -o out_ref2va.mp4 --max-time 14400 \
-X POST "http://127.0.0.1:8000/v1/videos/sync" \
-F 'prompt=傍晚的小厨房,夕阳余晖从窗边洒进来,旧木桌上放着洗到一半的马克杯和起雾的玻璃瓶,悬挂的抹布轻轻晃动。画面带有手持拍摄的轻微晃动和逆光曝光波动,生活感十足。环境音是安静的厨房底噪。' \
-F 'input_reference=@/path/to/reference.jpg;type=image/jpeg' \
-F 'audio_reference=</path/to/audio_reference.json' \
-F 'width=1024' -F 'height=576' -F 'fps=24' \
-F 'num_inference_steps=60' -F 'flow_shift=12' -F 'seed=1101' \
-F 'extra_params={"task":"ref2va","duration":5,"audio_flow_shift":3.0}'
```
`audio_reference` is a JSON object, which is why the example loads it from a file
with curl's `<` syntax. Its `audio_url` must be an `http(s)` URL or a data URL;
a bare filesystem path is rejected with `Invalid audio_reference.audio_url`.
```json
{"audio_url": "https://example.com/reference.wav"}
```
Reference images are validated before inference: the short edge must be at least
256 pixels, the long edge at most 5760, and the aspect ratio must fall between
0.4 and 2.5.
The response headers carry the measurements used in the table above:
```bash
cat headers_t2va.txt
ffmpeg -v error -i out_t2va.mp4 -f null - && echo DECODE_OK
```
+4 -67
View File
@@ -79,8 +79,7 @@ sharing the tokenizer, processor, text encoder, and VAEs from `FL2VA`.
### Multi-NPU: 768P combined-service configuration
Use Ulysses sequence parallelism degree 8, text-encoder tensor parallelism
degree 8, native tiled VAE patch parallelism degree 8, and distributed
layerwise offload:
degree 8, native tiled VAE patch parallelism degree 8, and layerwise offload:
```bash
export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
@@ -99,7 +98,7 @@ vllm serve "${MODEL}" \
--usp 8 \
--ring 1 \
--text-encoder-tp-size 8 \
--enable-distributed-layerwise-offload \
--enable-layerwise-offload \
--vae-parallel-mode tile \
--vae-use-tiling \
--vae-patch-parallel-size 8 \
@@ -113,11 +112,6 @@ H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
The same endpoint accepts `task=t2va`, `task=fl2va`, and `task=ref2va`; no
partition restart is required. Layerwise offload applies to both DiTs.
On Atlas 800I A3 (64 GB HBM per device) the combined service does not fit at
768P without offloading or sharding: use distributed layerwise offload (as
above) or HSDP — see
[§ Memory and attention optimizations](#memory-and-attention-optimizations-a3).
### Optional optimizations
Two independent optimizations may be enabled on top of the configuration
@@ -140,63 +134,6 @@ Keep `--ring 1` when using RainFusion: the `rf_v2` kernel ranks key blocks
over the whole sequence, so ring parallelism would split away the keys it
needs. Scale with `--usp` instead.
## Memory and attention optimizations (A3)
### Fitting 768P into 64 GB HBM: distributed layerwise offload or HSDP
Each Atlas 800I A3 NPU has 64 GB of HBM, which is not enough for the
combined MiniMax-H3 service at 768P. Enable **one** of the following at
server startup:
**Distributed layerwise offload** — DiT layers are offloaded with parameters
gathered across the parallel group instead of replicated per rank:
```bash
--enable-distributed-layerwise-offload
```
Measured peak NPU memory is about 45 GB per device for Ref2VA with a 13.88 s
input video generating a 15 s 768P (1344x768) output video. Host (CPU) memory
usage is high with this option.
**HSDP** — hybrid sharded data parallelism for the DiT parameters. The
multi-stream memory-reuse knob is mandatory with this flag:
```bash
export MULTI_STREAM_MEMORY_REUSE=2
```
```bash
--use-hsdp
```
Host memory usage is small, but large shapes may still OOM; HBM usage
optimizations for this path are ongoing.
### FLASH_ATTN backend with MindIE-SD
Keep `--diffusion-attention-backend FLASH_ATTN` and install MindIE-SD (see
[§ Environment](#environment)). On NPU this backend carries most of the
memory-reduction and performance work: a mask-free packed varlen path that
never materializes the quadratic `full_qk` padding mask, and K/V prefix
slicing, both driven by the packed `cu_seqlens` metadata emitted by the H3
transformer.
### Optional: LaserAttention fused kernel
For an additional attention speedup, select the Ascend LaserAttention fused
kernel before starting the server:
```bash
export MINDIE_SD_FA_TYPE="ascend_laser_attention"
```
This requires the FLASH_ATTN backend and MindIE-SD. H3 automatically applies
exact power-of-two input pre-scaling (`laser_input_scale=256`) so the
kernel's fp16 workspace cannot overflow on outlier activations. Measured
speedup numbers will be added here.
## HTTP API examples
The request format is identical to the GPU recipe; see
@@ -217,8 +154,8 @@ configuration above:
| Workload | Configuration |
|----------|---------------|
| T2VA, 209 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile |
| T2VA, 209 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
These measurements describe the validated shapes rather than a general
throughput guarantee.
@@ -1,346 +0,0 @@
# MiniMax-H3 on RTX PRO 6000 Blackwell GPUs
This recipe runs MiniMax-H3 in BF16 on 96 GiB RTX PRO 6000 Blackwell GPUs.
The additional HBM over the RTX PRO 5000 profile removes the four-GPU
minimum: two GPUs are enough to keep the model resident with TP2. Four and
eight GPUs raise the Ulysses degree to shard the attention sequence further.
CPU offload and distributed layerwise offload are not required in any of
these configurations.
Validated on:
- Host: YLX Y762
- GPUs: 8 × RTX PRO 6000 Blackwell (96 GiB)
- Device order: default (`CUDA_VISIBLE_DEVICES` not set)
- Driver / CUDA: 580.105.08 / 13.0
- Container image: `vllm/vllm-omni:minimax-h3`
- Workload: T2VA, 1344×768, `duration=5.0`, 50 steps, `flow_shift=12`, `seed=1101`
## Capacity requirements
| Resource | Two GPUs | Four GPUs | Eight GPUs |
| --- | ---: | ---: | ---: |
| GPU HBM | 96 GiB per GPU | 96 GiB per GPU | 96 GiB per GPU |
| Checkpoint storage | 135 GiB per partition | 135 GiB per partition | 135 GiB per partition |
| Available system RAM | 200 GiB minimum | 200 GiB minimum | 200 GiB minimum |
| Recommended system RAM | 384 GiB | 384 GiB | 384 GiB |
| Measured peak HBM per GPU | 77.49 GiB | 66.44 GiB | 61.07 GiB |
`FL2VA` and `Ref2VA` are separate checkpoint partitions. Start one server at
a time on a host sized for the minimum system-memory requirement, or pass
`--task-type fl2va` or `--task-type ref2va` to download and load a single
partition. The two-server layout described in the eight-GPU section doubles
both the storage and the host-memory requirement.
## Parallelism summary
| Configuration | TP | Ulysses | Text-encoder TP | VAE patch parallel |
| --- | ---: | ---: | ---: | ---: |
| Two GPUs | 2 | 1 | 2 | 2 |
| Four GPUs | 2 | 2 | 4 | 4 |
| Eight GPUs | 2 | 4 | 8 | 8 |
| Eight GPUs, headroom variant | 4 | 2 | 8 | 8 |
Per-GPU weight residency is set by the tensor-parallel degree alone. Raising
the Ulysses degree lowers activation memory and per-step latency but leaves
weight residency unchanged. Ring parallelism stays at 1, and H3 is
CFG-distilled, so `--cfg-parallel-size` must also remain 1.
## PCIe topology and GPU order
RTX PRO 6000 Blackwell does not provide NVLink. Before starting the server,
inspect the host topology:
```bash
nvidia-smi topo -m
nvidia-smi nvlink -s
```
Tensor-parallel groups are consecutive ranks. Ulysses groups are strided by
the tensor-parallel degree. Confirm this grouping against the server log
before tuning the device order on a new host.
For the two-GPU configuration, pick a single `PXB` pair on one NUMA node.
For the four-GPU configuration, pick two `PXB` pairs on the same node and
order them so that each Ulysses group lands inside one pair. If physical
GPUs `(0,1)` and `(2,3)` are the two local pairs, the order is
`CUDA_VISIBLE_DEVICES=0,2,1,3`. Do not copy these IDs blindly: reproduce the
same relationship on the target host.
Device-order pinning and NUMA binding are tuning knobs, not prerequisites.
The measurements in this recipe were taken without either, so any gain from
them is additional to the numbers reported below and is host-specific.
## Two-GPU serving configuration
Two 96 GiB GPUs hold the BF16 model with TP2 alone. There is no Ulysses
group, so the DiT attention sequence is not sharded and activation memory is
higher than in the four-GPU profile at the same output shape. Re-measure
peak memory before increasing the output shape, the reference-input length,
or concurrency.
```bash
export MODEL_ROOT=/path/to/MiniMax-H3
export MODEL="${MODEL_ROOT}/FL2VA"
export PORT=8000
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
vllm serve "${MODEL}" \
--omni \
--host 0.0.0.0 \
--port "${PORT}" \
--trust-remote-code \
--num-gpus 2 \
--tensor-parallel-size 2 \
--usp 1 \
--ring 1 \
--text-encoder-tp-size 2 \
--vae-patch-parallel-size 2 \
--vae-parallel-mode tile \
--vae-use-tiling \
--diffusion-attention-backend CUDNN_ATTN
```
## Four-GPU serving configuration
Four GPUs keep TP2 for the weight shard and add Ulysses2, which shards the
attention sequence and lowers per-GPU activation memory. Text-encoder TP4
and VAE patch parallelism 4 spread the Qwen3-VL encoder and tiled decode
across all four GPUs. TP1 is not an option on this card: an unsharded BF16
partition does not fit in 96 GiB.
```bash
export MODEL_ROOT=/path/to/MiniMax-H3
export MODEL="${MODEL_ROOT}/FL2VA"
export PORT=8000
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
vllm serve "${MODEL}" \
--omni \
--host 0.0.0.0 \
--port "${PORT}" \
--trust-remote-code \
--num-gpus 4 \
--tensor-parallel-size 2 \
--usp 2 \
--ring 1 \
--text-encoder-tp-size 4 \
--vae-patch-parallel-size 4 \
--vae-parallel-mode tile \
--vae-use-tiling \
--diffusion-attention-backend CUDNN_ATTN
```
## Eight-GPU serving configuration
Eight GPUs extend the Ulysses degree to 4 while keeping TP2. This does not
lower per-GPU weight residency relative to the four-GPU profile; it lowers
activation memory and per-step latency by sharding the attention sequence
four ways.
```bash
export MODEL_ROOT=/path/to/MiniMax-H3
export MODEL="${MODEL_ROOT}/FL2VA"
export PORT=8000
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
vllm serve "${MODEL}" \
--omni \
--host 0.0.0.0 \
--port "${PORT}" \
--trust-remote-code \
--num-gpus 8 \
--tensor-parallel-size 2 \
--usp 4 \
--ring 1 \
--text-encoder-tp-size 8 \
--vae-patch-parallel-size 8 \
--vae-parallel-mode tile \
--vae-use-tiling \
--diffusion-attention-backend CUDNN_ATTN
```
### Rank ordering across two sockets
With TP2 and Ulysses4 the two Ulysses groups are ranks `(0,2,4,6)` and
`(1,3,5,7)`, and the four tensor-parallel pairs are `(0,1)`, `(2,3)`,
`(4,5)`, and `(6,7)`. On a dual-socket host these two collectives cannot
both be socket-local. If you choose to pin the device order, the trade-off
is:
| Device order | Socket-local collective | Crosses sockets |
| --- | --- | --- |
| `0,4,1,5,2,6,3,7` | Ulysses all-to-all | TP all-reduce |
| `0,1,2,3,4,5,6,7` | TP all-reduce | Ulysses all-to-all |
The first order assumes physical GPUs `0-3` are on NUMA node 0 and `4-7` on
node 1. Reproduce that relationship on the target host rather than copying
the IDs. The reported measurements use the default order and no NUMA
binding; the two orders were not compared on this host.
### Headroom variant: TP4 with Ulysses2
Raising the tensor-parallel degree to 4 halves per-GPU weight residency and
leaves room for long Ref2VA references, larger output shapes, or
concurrency greater than one. It also adds a four-way all-reduce per DiT
block over a link without NVLink. Use it when the TP2 profile runs out of
headroom, not as the default. Replace the parallel flags above with:
```bash
--tensor-parallel-size 4 \
--usp 2 \
```
### Two independent four-GPU servers
An eight-GPU host can instead run one FL2VA server and one Ref2VA server
side by side, each pinned to its own NUMA node with the four-GPU
configuration. This removes every cross-socket collective and lifts the
one-server-at-a-time restriction, at the cost of doubling the requirements
to 400 GiB available system RAM and 270 GiB of model storage.
Use `CUDA_VISIBLE_DEVICES=0,2,1,3`, `--cpunodebind=0 --membind=0`, and
`--port 8091` for the FL2VA server. Use `CUDA_VISIBLE_DEVICES=4,6,5,7`,
`--cpunodebind=1 --membind=1`, and `--port 8092` for the Ref2VA server.
This layout is described for completeness and was not part of the measured
runs.
## Shared serving notes
Do not add `--enforce-eager` for a performance run. Warm the server before
measuring so regional compilation is outside the measured request. On this
host the first request after startup ran 19% slower than the steady state
(107.58 s versus 90.48 s on eight GPUs); two warmup requests were enough to
converge.
For Ref2VA on a single-server layout, stop the FL2VA server and restart the
same command with `MODEL="${MODEL_ROOT}/Ref2VA"`.
## Attention backend
RTX PRO 6000 Blackwell is an SM120 part, so the datacenter-Blackwell
`TRTLLM_ATTN` default does not apply here. Every configuration above selects
cuDNN BF16 attention explicitly. This keeps the recipe independent of
platform-default backend changes and does not depend on the experimental
SM120 kernel.
## Optional: online FP8
`--quantization fp8` quantizes eligible DiT linears at load time and is
compatible with tensor parallelism and VAE tiling. Use it when the resident
BF16 peak leaves too little headroom for long Ref2VA references, a larger
output shape, or concurrency greater than one. It cannot be combined with
layerwise offload.
## Target-hardware validation
T2VA, 1344×768, `duration=5.0`, `fps=24`, 50 steps, `flow_shift=12`,
`seed=1101`, BF16, `CUDNN_ATTN`, tiled VAE, one request at a time. Servers
were started with the commands above — default device order, no NUMA
binding — plus `--enable-diffusion-pipeline-profiler`. Two warmup requests
preceded each measured request.
| Measurement | Two GPUs | Four GPUs | Eight GPUs |
| --- | ---: | ---: | ---: |
| Topology | TP2 × USP1 | TP2 × USP2 | TP2 × USP4 |
| Text encode | 0.040 s | 0.044 s | 0.035 s |
| Denoise, 50 steps | 278.55 s | 168.73 s | 87.90 s |
| Per step | 5.571 s | 3.375 s | 1.758 s |
| VAE decode | 5.396 s | 2.791 s | 1.798 s |
| Client E2E | 284.76 s | 172.32 s | 90.48 s |
| Peak HBM per GPU | 77.49 GiB | 66.44 GiB | 61.07 GiB |
| Headroom below 96 GiB | 18.5 GiB | 29.6 GiB | 34.9 GiB |
Stage times are read from the `X-Stage-Durations` response header of
`/v1/videos/sync`. Peak memory is the maximum of `nvidia-smi
--query-gpu=memory.used` sampled at 1 Hz across every device for the
duration of the measured request. Per step is denoise wall time divided by
the 50 requested steps. Stage times sum to roughly 0.75 s less than
end-to-end in all three configurations; queueing, result transfer, and MP4
muxing sit outside the profiled stages.
### Scaling behaviour
Denoise speedup is not uniform across the sweep: 1.65× from two to four
GPUs (83% efficiency) but 1.92× from four to eight (96%). Two GPUs run
Ulysses1, so there is no sequence parallelism at all; going to four GPUs
pays the one-time cost of introducing the all-to-all, while going from four
to eight only widens an all-to-all that already exists. Expect the second
doubling to pay off better than the first on this topology.
VAE decode behaves in the opposite direction — 1.93× from two to four GPUs
but only 1.55× from four to eight — as patch parallelism produces smaller
tiles and tile-boundary overhead grows.
Text encoding is 0.04 s in every configuration and is not worth optimising;
denoise is 96-98% of end-to-end throughout.
### Per-GPU memory model
Peak memory falls by 11.05 GiB from two to four GPUs and by a further
5.37 GiB from four to eight — halving each time the Ulysses degree doubles,
because all three configurations shard the DiT with TP2 and hold identical
weights per GPU. Fitting the first two points gives:
```
peak HBM per GPU ≈ 55.4 GiB + 22.1 GiB / ulysses_degree
```
That fit predicts 60.9 GiB at Ulysses4 against 61.07 GiB measured, a 0.25%
error. The practical consequence is a floor near 55 GiB per GPU: adding
GPUs past eight will not reduce the per-GPU requirement further. To serve
on smaller cards, raise the tensor-parallel degree instead.
Re-measure memory for longer reference inputs, concurrency greater than
one, or a different output shape.
### Not yet measured
Ref2VA latency and memory have not been measured on this host. Ref2VA
drives the Qwen3-VL presentation from tens of tokens for text-only prompts
to several thousand with image or video references, so T2VA numbers do not
transfer to it. The eight-GPU device-order comparison in the section above
was also not run.
## T2VA request example
`t2va` requires an explicit `aspect_ratio` even when `width` and `height`
are supplied.
```bash
export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync"
curl -sS --max-time 1800 -X POST "${API_URL}" \
-F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \
-F 'width=1344' \
-F 'height=768' \
-F 'aspect_ratio=16:9' \
-F 'fps=24' \
-F 'num_inference_steps=50' \
-F 'flow_shift=12' \
-F 'seed=1101' \
-F 'extra_params={"task":"t2va","duration":5.0,"audio_flow_shift":3.0}' \
-o t2va.mp4
```
To collect the stage breakdown and peak memory for a request, add
`--enable-diffusion-pipeline-profiler` to the server command and `-D
headers.txt` to the curl invocation, then read `X-Stage-Durations` and
`X-Peak-Memory-MB` from the saved headers.
## Ref2VA request example
`ref2va` defaults to a 16:9 aspect ratio, so `aspect_ratio` is optional
here.
```bash
curl -X POST "http://127.0.0.1:${PORT}/v1/videos/sync" \
--fail-with-body -w '\nHTTP %{http_code}\n' \
--max-time 1200 \
-F "input_reference=@/root/hand.jpg;type=image/jpeg" \
-F "audio_reference=</root/audio_ref.json" \
-F "prompt=2D动画融合在一起的影像。夕阳余晖残留在窗边,生活感十足的小厨房里有旧木桌、洗到一半的马克杯、起雾的玻璃瓶、悬挂的抹布。画面带有智能手机单手拍摄的手抖、近距离对焦的犹豫、逆光曝光波动。要像在家中慌忙拍下某个不可思议事件的自然质感,不要广告影像的精心整理。声音只用厨房环境声与手绘生物柔和的电子音、小小的叫声。" \
-F 'width=1344' -F 'height=768' -F 'fps=24' \
-F 'num_inference_steps=60' -F 'flow_shift=12' -F 'seed=1101' \
-F 'extra_params={"task":"ref2va","duration":8,"audio_flow_shift":3.0}' \
-o /root/out_ref2va.mp4
```
-262
View File
@@ -1,262 +0,0 @@
# MiniMax-H3 on DGX Spark (GB10)
This recipe uses online FP8 weight quantization, tiled VAE decode, and a single
resident partition. GB10 is a unified-memory platform, so unlike the discrete-GPU
recipes it uses **no offload of any kind** — see the capacity note below.
Validated on:
- Host: DGX Spark (GB10), aarch64
- GPUs: 1 (unified memory)
- Driver: <FILL_DRIVER_VERSION>
- vLLM: 0.26.0
- vLLM-Omni: `main` at `e1aa6eae75c460cd1893bc320546e81e66973831`
- Workloads: T2VA and Ref2VA, 960x576, `duration=8.0`, 50 steps, `flow_shift=12`,
`seed=1101`
## Capacity requirements
| Resource | DGX Spark (GB10) |
| --- | ---: |
| Unified memory | 128 GB LPDDR5X (~121 GiB usable) |
| GPUs | 1 |
| Checkpoint storage | 135 GiB per partition |
| Measured peak allocator high-water, T2VA (FP8) | 97.7 GiB |
| Measured peak allocator high-water, Ref2VA (FP8) | 102.8 GiB |
`FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one server
at a time.
On GB10 the CPU and GPU share one physical memory pool, which changes the capacity
math relative to the RTX and datacenter recipes:
- **Do not use `--enable-distributed-layerwise-offload`.** DLO stages rank-local
weights into pinned host memory, which on GB10 is the same pool the weights are
already in. Peak usage roughly doubles and the process is killed by the OOM
killer (`Exit code: -9`) shortly after `Enabling offloader backend`.
- **Do not use `--enable-cpu-offload`.** Moving weights from "VRAM" to "host RAM"
frees nothing here.
- **`--quantization fp8` is mandatory.** A BF16 partition is 135 GiB and does not
fit in 121 GiB. Online FP8 quantizes the DiT only (62 GiB to ~31 GiB); the
Qwen3-VL text encoder and both VAEs stay BF16.
With ~97.7 GiB of the pool held by the allocator at peak, T2VA leaves roughly
23 GiB of headroom for the OS, page cache, and any other process on the box.
Ref2VA is tighter on both counts. The allocator high-water rises to 102.8 GiB,
and the allocator figure understates true pool usage.
Sampling `/proc/meminfo` at 1 Hz across a 50-step Ref2VA request on a host with
`MemTotal: 127600524 kB` (121.7 GiB), `MemAvailable` started at 8.1 GiB, fell to
a trough of **7.4 GiB**, and recovered to 8.9 GiB. At the trough the pool held
roughly **114.3 GiB** — 11.4 GiB more than the 102.8 GiB the allocator reports,
which is the CUDA context, non-PyTorch allocations, and page cache the header
does not count.
Treat this platform as strictly single-tenant for Ref2VA. There is no room for a
second process, and adding reference images, reference videos, or a larger output
shape will run into the OOM killer.
## One GB10: 960x576, 8 seconds
```bash
CUDA_VISIBLE_DEVICES=0 \
FLASHINFER_DISABLE_VERSION_CHECK=1 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=7200 \
vllm serve /path/to/MiniMax-H3/FL2VA \
--omni --trust-remote-code --host 0.0.0.0 --port 8000 \
--init-timeout 3600 \
--num-gpus 1 --tensor-parallel-size 1 --text-encoder-tp-size 1 \
--usp 1 --ring 1 --vae-patch-parallel-size 1 \
--vae-parallel-mode tile --vae-use-tiling \
--quantization fp8 --enforce-eager \
--diffusion-attention-backend CUDNN_ATTN
```
Wait for `Application startup complete`, then run a short smoke test before
committing to a full 50-step request.
`t2va` requires an explicit `aspect_ratio` even when `width` and `height` are
supplied; omitting it fails the request with
`OmniClientError: t2va requires an explicit aspect_ratio`.
```bash
curl --fail-with-body -sS -X POST http://127.0.0.1:8000/v1/videos/sync \
-F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \
-F 'width=960' \
-F 'height=576' \
-F 'aspect_ratio=16:9' \
-F 'fps=24' \
-F 'num_inference_steps=10' \
-F 'flow_shift=12' \
-F 'seed=1101' \
-F 'extra_params={"task":"t2va","duration":4.0,"audio_flow_shift":3.0}' \
-o h3-gb10-smoke.mp4
```
The full-quality request uses 50 steps:
```bash
curl --fail-with-body -sS -X POST http://127.0.0.1:8000/v1/videos/sync \
-F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \
-F 'width=960' \
-F 'height=576' \
-F 'aspect_ratio=16:9' \
-F 'fps=24' \
-F 'num_inference_steps=50' \
-F 'flow_shift=12' \
-F 'seed=1101' \
-F 'extra_params={"task":"t2va","duration":8.0,"audio_flow_shift":3.0}' \
-o h3-gb10-t2va.mp4
ffprobe -v error -show_entries \
stream=index,codec_name,width,height,r_frame_rate,sample_rate,channels \
-of json h3-gb10-t2va.mp4
```
Both requests return `200 OK` with an MP4 body: H.264 video at the requested
geometry plus a 32 kHz stereo AAC track.
## One GB10: Ref2VA, 960x576, 8 seconds
`FL2VA` and `Ref2VA` are separate 135 GiB partitions and only one fits at a time.
Stop the FL2VA server before starting this one.
```bash
CUDA_VISIBLE_DEVICES=0 \
FLASHINFER_DISABLE_VERSION_CHECK=1 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
vllm serve /path/to/MiniMax-H3/Ref2VA \
--omni --trust-remote-code --host 0.0.0.0 --port 8000 \
--init-timeout 3600 \
--num-gpus 1 --tensor-parallel-size 1 --text-encoder-tp-size 1 \
--usp 1 --ring 1 --vae-patch-parallel-size 1 \
--vae-parallel-mode tile --vae-use-tiling \
--quantization fp8 --enforce-eager \
--diffusion-attention-backend CUDNN_ATTN
```
The 7200 s timeout used for T2VA is too short here: a 50-step Ref2VA request
takes over an hour on this platform. Use 14400 s.
Unlike `t2va`, `ref2va` defaults to a 16:9 output ratio, so `aspect_ratio` can be
omitted when `width` and `height` are supplied.
```bash
curl --fail-with-body -sS -X POST http://127.0.0.1:8000/v1/videos/sync \
-D ref2va-headers.txt \
-w 'client_e2e_s=%{time_total}\n' \
--max-time 14400 \
-F 'input_reference=@/path/to/reference.jpg;type=image/jpeg' \
-F 'prompt=A hand rests on a wooden table in warm evening light, then slowly opens and turns toward the camera, with quiet room ambience.' \
-F 'width=960' \
-F 'height=576' \
-F 'fps=24' \
-F 'num_inference_steps=50' \
-F 'flow_shift=12' \
-F 'seed=1101' \
-F 'extra_params={"task":"ref2va","duration":8.0,"audio_flow_shift":3.0}' \
-o h3-gb10-ref2va.mp4
```
The response is `200 OK` with an MP4 body: H.264 at the requested geometry plus a
32 kHz stereo AAC track. `ffprobe` reports 8.032 s for `duration=8.0`.
## Measured latency
Single GB10, FP8, eager, cuDNN attention, tiled VAE, one request at a time:
| Workload | Steps | Text encode | Denoise | Per step | VAE decode | MP4 mux | Client E2E | Peak memory |
| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
| T2VA, 4 s, smoke shape | 10 | — | — | 8.68 s | — | 0.49 s | 95.9 s | — |
| T2VA, 8 s, 960x576 | 50 | 0.25 s | 2088 s | 42.61 s | 70 s | 4.84 s | 2169.4 s (36 min 9 s) | 97.7 GiB |
| Ref2VA, 8 s, 960x576, one image (3 runs) | 10 | 8.43-10.94 s | 665.5-748.8 s | 66.55-74.88 s | 67.9-69.8 s | 3.63 s | 770.3-861.1 s | 102.8 GiB |
| Ref2VA, 8 s, 960x576, one image | 50 | 9.66 s | 4039.6 s | 80.79 s | 69.4 s | 3.63 s | 4157.0 s (69 min 17 s) | 102.8 GiB |
How these were measured:
- Stage times come from `--enable-diffusion-pipeline-profiler`, read from the
`X-Stage-Durations` response header of `/v1/videos/sync`. The keys are
`MiniMaxH3Pipeline.encode_prompt` (text encode),
`MiniMaxH3Pipeline.diffuse` (denoise), and `MiniMaxH3Pipeline.decode`
(VAE decode).
- MP4 mux is not a header field. It is the
`Video response encoding (MP4 bytes)` line in the server log.
- Peak memory is the `X-Peak-Memory-MB` header, which reports the CUDA allocator
high-water mark (`torch.cuda.max_memory_reserved`) for the request. It excludes
the CUDA context and non-PyTorch allocations, so true pool usage is somewhat
higher. `nvidia-smi --query-gpu=memory.used` is not usable on GB10 — unified
memory has no discrete VRAM to report and the field returns `[N/A]`. See the
capacity section for a `/proc/meminfo` cross-check.
- For T2VA, text encode and VAE decode were measured on a separate 10-step run at
the same 960x576 / 8 s geometry. Both are independent of step count, so they
transfer to the 50-step row unchanged. The Ref2VA MP4 mux figure transfers the
same way, from the 10-step run at identical output geometry and duration.
- Ref2VA peak memory is byte-identical between the 10-step and 50-step runs
(`X-Peak-Memory-MB: 105318.000` in both), confirming that peak is set by output
geometry and reference input rather than step count.
- Stage times do not sum exactly to end-to-end. Queueing, result transfer, and
engine bookkeeping sit outside the profiled stages.
Denoise dominates at 96% of end-to-end, so step count and geometry are the only
knobs that materially change the wall clock on this platform.
Ref2VA costs 1.92x the wall clock of T2VA at the same geometry and step count:
4157.0 s versus 2169.4 s. Text encoding shows the most dramatic ratio — 9.66 s
versus 0.25 s, 39x — but at 0.2% of end-to-end it is not what makes Ref2VA slow.
The cost is in denoise. A single reference image expands the Qwen3-VL
presentation to 6883 tokens (logged as `MiniMax H3 ref2va Qwen presentation:
6883 tokens, 1 reference images`), and those tokens sit in the attention context
of every denoise step. Per-step time rises from 42.61 s to 80.79 s, and that
1.90x multiplies across all 50 steps. VAE decode is unchanged at ~69 s, as
expected for an identical output shape.
Repeated identical 10-step T2VA requests varied by up to 23% in denoise time
(397 s to 490 s) on a machine that had already been under sustained load. Treat
these as single-run measurements on one machine, not a throughput guarantee, and
expect the slower end of that range for back-to-back jobs.
Ref2VA per-step time is load-dependent in the same way. On an idle machine the
first 10-step request ran at 66.55 s per step. The 50-step request that followed
occupied the box for 69 consecutive minutes and averaged 80.79 s per step. Two
further 10-step requests taken immediately afterwards settled at 74.88 s and
74.31 s — 12.5% slower than the idle run, and within 0.8% of each other, so this
is a stable thermal state rather than measurement noise.
Plan for 81 s per step when sizing a full 50-step Ref2VA job; 66 s per step is
only reachable on a cold machine and does not hold for the length of one request.
Ref2VA text encode converged downward across the four runs (10.94 s, 9.66 s,
8.52 s, 8.43 s) as first-touch costs fell away; treat ~8.5 s as steady state.
VAE decode was stable at 67.9-69.8 s throughout.
Set `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` well above the expected request time — the
default 1800 s is shorter than a 50-step run on this platform, and 7200 s is
shorter than a 50-step Ref2VA run.
## Known limitations
- One GPU means no Ulysses, ring, TP, or VAE patch parallelism. Pass
`--usp 1 --ring 1 --vae-patch-parallel-size 1` explicitly if a profile or
script would otherwise supply higher degrees.
- `--enforce-eager` is used because regional compile adds startup time without a
parallelism win on a single rank.
- `t2va` rejects requests without an explicit `aspect_ratio`. When `width` and
`height` are both supplied they take precedence and `aspect_ratio` only has to
name one of the supported ratios.
- Sustained load degrades throughput. See the variance notes above.
- Only one partition is resident. For Ref2VA, stop the FL2VA server and restart
with `/path/to/MiniMax-H3/Ref2VA`. Ref2VA with one reference image is measured
above; reference videos and additional images push the token count higher and
were not measured. With `MemAvailable` bottoming out at 7.4 GiB during a
single-image request, there is very little room to grow — re-measure before
adding references or raising the output shape.
- Ref2VA requires a vLLM-Omni build newer than `v0.26.0`. See the version note
under **Validated on**: FP8 weight loading and image-only Ref2VA are both
broken on `release/v0.26.0`.suggestion is `v0.26.1`.
- Online FP8 is incompatible with layerwise offload — the offload path produces a
weight stride the Cutlass FP8 kernel rejects. This is not a practical
restriction here since offload is unusable on GB10 anyway.
```
-69
View File
@@ -750,73 +750,8 @@ audio spectral cosine 0.9589 (minimum 0.80), and audio RMS ratio 0.9342. The
resident per-GPU peak was 68.52 GiB for BF16 and 53.51 GiB for FP8, a 22%
reduction.
## TeaCache acceleration
TeaCache reuses DiT block residuals across denoising steps when consecutive
timestep embeddings are similar. MiniMax-H3 TeaCache is currently calibrated
only for the FL2VA partition. In combined serving, FL2VA requests use TeaCache
while Ref2VA requests run uncached; Ref2VA-only serving rejects TeaCache.
TeaCache and Cache-DiT are mutually exclusive; pick one cache backend per server.
The model-specific default and examples use `rel_l1_thresh=0.17`, which
provided the best conservative speed/quality balance in the validated 107-frame
T2VA workload. Lower values
may produce few or no cache hits, while higher values can improve performance
at the cost of output quality. Validate the threshold on representative
prompts and generation settings before changing it.
### Offline (Python API)
Export the directory containing the `FL2VA` and `Ref2VA` subdirectories before
running the Python example, for example `export MODEL_ROOT=/models/MiniMax-H3`.
```python
import os
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model=os.path.join(os.environ["MODEL_ROOT"], "FL2VA"),
cache_backend="tea_cache",
cache_config={"rel_l1_thresh": 0.17},
trust_remote_code=True,
enable_cpu_offload=True,
)
outputs = omni.generate(
"A quiet cinematic night scene with matching ambient sound.",
OmniDiffusionSamplingParams(
height=256,
width=448,
num_frames=29,
fps=24,
num_inference_steps=50,
seed=42,
extra_args={
"task": "t2va",
"duration": 4.0,
"aspect_ratio": "16:9",
"flow_shift": 12.0,
"audio_flow_shift": 3.0,
},
),
)
```
### Online serving
```bash
vllm serve "${MODEL_ROOT}/FL2VA" \
--omni \
--trust-remote-code \
--cache-backend tea_cache \
--cache-config '{"rel_l1_thresh":0.17}'
```
## Known limitations
- TeaCache is calibrated for FL2VA only; Ref2VA requests run uncached.
- Combined serving requires sibling `FL2VA` and `Ref2VA` directories, loads
both task-specific DiTs, and loads shared components once from `FL2VA`.
- H3 currently executes one generation request per diffusion batch.
@@ -835,10 +770,6 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
- A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
pure Ulysses.
- Online FP8 is not compatible with layerwise offload.
- TeaCache and Cache-DiT cannot be enabled on the same server.
- Image+audio Ref2VA accepts exactly one image and one audio reference.
- Video Ref2VA accepts one or more video files, but not an additional standalone
audio reference.
- Pure Ulysses still replicates the full DiT on every rank, so smaller-memory GPUs
cannot use `--usp N --tp 1` as a resident capacity path. Use DiT tensor parallelism
or model-level CPU offload; text-encoder TP alone is not sufficient.
+1 -2
View File
@@ -30,9 +30,8 @@ also provided.
`hf_config.version="4.5"`):
- Default single-GPU compatibility layout (auto-loaded):
[`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml)
- Recommended 2-GPU continuous-batching layout:
- 2-GPU and 3-GPU layouts:
[`vllm_omni/deploy/minicpmo_4_5_2gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_2gpu.yaml),
- 3-GPU layout:
[`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml)
- 8x RTX 4090 layout:
[`vllm_omni/deploy/minicpmo_4_5_8x4090.yaml`](../../vllm_omni/deploy/minicpmo_4_5_8x4090.yaml)
-6
View File
@@ -19,11 +19,6 @@ add files for this repo. Use one Markdown file per model family by default.
> Individual pages resolve at `https://recipes.vllm.ai/<provider>/<model-id>`,
> e.g. [`bosonai/higgs-audio-v3-tts-4b`](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b).
The `Recipe` column in the Supported Models documentation links to the
published recipe when available and otherwise to the corresponding recipe in
this directory. The recipe files remain the source of truth for commands and
validation details.
Example layout:
```text
@@ -55,7 +50,6 @@ recipes/
| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
| [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
| [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB |
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
| [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) |
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
@@ -1,282 +0,0 @@
# LongCat-Video-Avatar-1.5
> Audio-driven avatar video generation (AT2V / AI2V)
## Summary
- Vendor: meituan-longcat
- Model: `meituan-longcat/LongCat-Video-Avatar-1.5`
- Task: Audio-to-video (AT2V) and audio-and-image-to-video (AI2V) avatar generation
- Mode: Offline inference
- Maintainer: Community
## When to use this recipe
Use this recipe to drive a talking or singing avatar from a speech clip, either
from text alone (AT2V) or from a reference image (AI2V), with the native
LongCat-Video-Avatar-1.5 pipeline in vLLM-Omni. It covers:
1. **Single-speaker AT2V / AI2V** — one 93-frame segment from one audio track.
2. **AVC continuation** — chained segments that cover a longer audio track.
3. **Multi-speaker AI2V** — several audio tracks placed on a reference image
through speaker bounding boxes.
The example saves an mp4 with the generated frames muxed with the input audio.
## References
- Upstream model card: <https://huggingface.co/meituan-longcat/LongCat-Video-Avatar-1.5>
- Upstream repository and example assets: <https://github.com/meituan-longcat/LongCat-Video>
- Related example under `examples/`: [`longcat_video/end2end.py`](../../examples/offline_inference/longcat_video/end2end.py)
## Hardware Support
## CUDA
### 1× NVIDIA H100 (80 GB)
#### Environment
- OS: Linux
- Python: 3.10+
- Driver: NVIDIA driver with CUDA 12.x
- vLLM version: Match the repository requirements for your checkout
- vLLM-Omni version or commit: Use the commit you are deploying from
#### Prerequisites
Install the LongCat Avatar optional runtime dependencies:
```bash
pip install "vllm-omni[longcat-video-avatar]"
```
The example uses the `imageio_ffmpeg` executable installed by vLLM-Omni's
`imageio[ffmpeg]` dependency.
#### Command
Audio-to-video (AT2V):
```bash
python examples/offline_inference/longcat_video/end2end.py \
--model meituan-longcat/LongCat-Video-Avatar-1.5 \
--stage at2v \
--audio /path/to/speech.wav \
--prompt "A person speaks calmly while facing the camera." \
--num-frames 93 \
--num-inference-steps 8 \
--output longcat_avatar_at2v.mp4
```
Audio-and-image-to-video (AI2V):
```bash
python examples/offline_inference/longcat_video/end2end.py \
--model meituan-longcat/LongCat-Video-Avatar-1.5 \
--stage ai2v \
--audio /path/to/speech.wav \
--image /path/to/reference.jpg \
--prompt "A person speaks calmly while facing the camera." \
--num-frames 93 \
--num-inference-steps 8 \
--output longcat_avatar_ai2v.mp4
```
#### Official Example Assets
Download the official Avatar example assets:
```bash
mkdir -p longcat_avatar_assets/assets/avatar/single longcat_avatar_assets/assets/avatar/multi
curl -L -o longcat_avatar_assets/assets/avatar/single/man.mp3 \
https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/single/man.mp3
curl -L -o longcat_avatar_assets/assets/avatar/single/man.png \
https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/single/man.png
curl -L -o longcat_avatar_assets/assets/avatar/single_example_1.json \
https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/single_example_1.json
curl -L -o longcat_avatar_assets/assets/avatar/multi/sing.png \
https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi/sing.png
curl -L -o longcat_avatar_assets/assets/avatar/multi/sing_man.WAV \
https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi/sing_man.WAV
curl -L -o longcat_avatar_assets/assets/avatar/multi/sing_woman.WAV \
https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi/sing_woman.WAV
curl -L -o longcat_avatar_assets/assets/avatar/multi/introduce.png \
https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi/introduce.png
curl -L -o longcat_avatar_assets/assets/avatar/multi/introduce_man.mp3 \
https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi/introduce_man.mp3
curl -L -o longcat_avatar_assets/assets/avatar/multi/introduce_woman.mp3 \
https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi/introduce_woman.mp3
curl -L -o longcat_avatar_assets/assets/avatar/multi_example_1.json \
https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi_example_1.json
curl -L -o longcat_avatar_assets/assets/avatar/multi_example_2.json \
https://raw.githubusercontent.com/meituan-longcat/LongCat-Video/main/assets/avatar/multi_example_2.json
export LONGCAT_VIDEO_ASSET_DIR="$PWD/longcat_avatar_assets/assets/avatar"
export AVATAR_MODEL=meituan-longcat/LongCat-Video-Avatar-1.5
export PROMPT="$(python -c 'import json, os; print(json.load(open(os.environ["LONGCAT_VIDEO_ASSET_DIR"] + "/single_example_1.json"))["prompt"])')"
```
`--base-model-dir` points at a local LongCat-Video base checkpoint for the
tokenizer, text encoder and VAE; omit it to download those components from
Hugging Face.
```bash
export BASE_MODEL_DIR=/cache/longcat_models/LongCat-Video
```
#### Official Asset Smoke Test
Audio-to-video with the official audio:
```bash
python examples/offline_inference/longcat_video/end2end.py \
--model "$AVATAR_MODEL" \
--base-model-dir "$BASE_MODEL_DIR" \
--stage at2v \
--audio "$LONGCAT_VIDEO_ASSET_DIR/single/man.mp3" \
--prompt "$PROMPT" \
--num-frames 93 \
--num-inference-steps 8 \
--output official_asset_longcat_avatar_at2v.mp4
```
Audio-and-image-to-video with the official audio and image:
```bash
python examples/offline_inference/longcat_video/end2end.py \
--model "$AVATAR_MODEL" \
--base-model-dir "$BASE_MODEL_DIR" \
--stage ai2v \
--audio "$LONGCAT_VIDEO_ASSET_DIR/single/man.mp3" \
--image "$LONGCAT_VIDEO_ASSET_DIR/single/man.png" \
--prompt "$PROMPT" \
--num-frames 93 \
--num-inference-steps 8 \
--output official_asset_longcat_avatar_ai2v.mp4
```
#### Official Single-Speaker AVC Continuation
`--num-segments` chains AVC segments; `auto` covers the full input audio.
```bash
python examples/offline_inference/longcat_video/end2end.py \
--model "$AVATAR_MODEL" \
--base-model-dir "$BASE_MODEL_DIR" \
--input-json "$LONGCAT_VIDEO_ASSET_DIR/single_example_1.json" \
--num-frames 93 \
--num-cond-frames 13 \
--num-segments 5 \
--num-inference-steps 8 \
--output official_asset_longcat_avatar_single_example_1_ai2v_5seg.mp4
```
Replace `--num-segments 5` with `--num-segments auto` to run until the full
official audio is covered. That is a slow archival validation path rather than a
regular smoke test: about 24m28s on a Modal H100 for the full
`single_example_1.json` audio with 8 inference steps.
#### Official Multi-Speaker AI2V
The multi-speaker cases use several audio tracks with speaker masks and optional
bounding boxes. The command below generates the first 93-frame segment, matching
the first call to the official `generate_ai2v()`.
```bash
python examples/offline_inference/longcat_video/end2end.py \
--model "$AVATAR_MODEL" \
--base-model-dir "$BASE_MODEL_DIR" \
--input-json "$LONGCAT_VIDEO_ASSET_DIR/multi_example_1.json" \
--num-frames 93 \
--num-inference-steps 8 \
--output official_asset_longcat_avatar_multi_example_1_ai2v.mp4
```
Use `multi_example_2.json` for the sequential-speaking case, and add
`--num-cond-frames 13 --num-segments auto` to cover the full audio with AVC
continuation.
#### Verification
```bash
ffprobe -hide_banner longcat_avatar_ai2v.mp4
```
The output must report one video stream and one audio stream: the example muxes
the generated frames with the input audio.
#### Notes
**Support matrix**
| Mode | Speakers | Image required | AVC continuation |
| --- | --- | --- | --- |
| AT2V | Single | No | Yes |
| AI2V | Single | Yes | Yes |
| AT2V | Multi | Not supported | Not supported |
| AI2V | Multi | Yes | Yes |
Multi-speaker Avatar requires AI2V because speaker bounding boxes and masks are
defined on a reference image; AT2V with several tracks fails with an explicit
error. AVC continuation uses full-audio embeddings only when `--num-segments` is
greater than `1` or set to `auto`; single-segment AT2V/AI2V keeps the shorter
first-segment audio path.
**Sampling parameters**
| Parameter | LongCat pipeline value | Description |
| --- | --- | --- |
| `--num-frames` | `num_frames` | Frames generated per segment. The default is 93 frames, or about 3.72 seconds at 25 FPS. LongCat normalizes other values to the temporal-VAE-compatible `4k+1` form. |
| `--num-inference-steps` | `steps` | Denoising steps. The distilled Avatar path uses 8 steps. |
| `--fps` | `save_fps` | Output frame rate, and the rate used to align audio embeddings with video frames. The default is 25 FPS. |
| `sampling_params.guidance_scale` | `text_guidance_scale` | Classifier-free guidance applied to the text prompt. |
| `sampling_params.guidance_scale_2` | `audio_guidance_scale` | Classifier-free guidance applied to the speech audio. |
With the distilled LoRA enabled (the default), the pipeline uses 8 inference
steps and sets both guidance scales to `1.0`. With `--no-use-distill`, the
example sets both to `4.0`. For AVC continuation every segment holds
`num_frames` frames, but each segment after the first reuses `num_cond_frames`
conditioning frames, so the unique frame count is
`num_frames + (num_segments - 1) * (num_frames - num_cond_frames)`.
**Key flags**
- `--stage`: `at2v` for audio-to-video, `ai2v` for audio-and-image-to-video.
- `--resolution`: `480p` or `720p`.
- `--input-json`: official LongCat Avatar JSON case, used for the single-speaker
and multi-speaker AI2V examples.
- `--num-segments`: number of AVC segments, or `auto` to cover the full audio.
- `--num-cond-frames`: previous-frame conditioning window for AVC continuation.
- `--use-int8` / `--no-use-int8`: load the official INT8 Avatar DiT weights by
default, or full precision weights.
- `--use-distill` / `--no-use-distill`: enable the official distilled LoRA path
by default.
- `--build-components-on-gpu`: build large Avatar components directly on GPU for
faster startup, at a higher peak VRAM cost. Disabled by default.
- `--base-model-dir`: local LongCat-Video base model directory for the
tokenizer, text encoder and VAE components.
**Memory usage**
Measured on a Modal H100 running the official 93-frame AI2V example with the
INT8 DiT, distilled LoRA and 8 inference steps:
| Component build | CLI option | Peak GPU memory | Model loading time |
| --- | --- | ---: | ---: |
| CPU (default) | None | About 41.0 GiB | About 130s |
| GPU | `--build-components-on-gpu` | About 56.8 GiB | About 33s |
The reported peak covers the end-to-end run. Building components on CPU only
reduces initialization memory; denoising still runs on GPU. These numbers are
memory guidance rather than a performance benchmark.
**Known limitations**
- Single GPU only. Cache acceleration, sequence/CFG/tensor parallelism, HSDP,
CPU offload, VAE patch parallelism and quantization are not wired up for this
pipeline yet; see the diffusion feature matrix in
[`docs/user_guide/diffusion_features.md`](../../docs/user_guide/diffusion_features.md).
- Offline inference only. Online serving is not supported yet.
+1 -159
View File
@@ -6,21 +6,13 @@ Unit tests for patch.py
"""
import asyncio
import base64
import json
import time
from types import SimpleNamespace
import pytest
from pytest_mock import MockerFixture
from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput
from vllm_omni.benchmarks.patch.patch import (
MixRequestFuncOutput,
async_request_openai_chat_omni_completions,
async_request_openai_realtime_duplex,
)
from vllm_omni.experimental.fullduplex.client import RealtimeEventCollector
from vllm_omni.benchmarks.patch.patch import MixRequestFuncOutput, async_request_openai_chat_omni_completions
pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu]
@@ -48,156 +40,6 @@ class MockResponse:
pass
@pytest.mark.asyncio
async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch):
class FakeRealtimeClient:
last_instance = None
def __init__(self, url):
assert url == "ws://localhost:8000/v1/realtime?duplex=1"
self.events = RealtimeEventCollector()
self.configure_kwargs = None
self.sent = []
self.response_count = 0
self.ack_count = 0
type(self).last_instance = self
async def __aenter__(self):
return self
async def __aexit__(self, *_args):
return None
async def configure(self, model, **kwargs):
assert model == "openbmb/MiniCPM-o-4_5"
self.configure_kwargs = kwargs
async def send(self, event):
self.sent.append(event)
if event["type"] != "response.create":
return
self.response_count += 1
response_id = f"resp-{self.response_count}"
now = time.monotonic()
self.events.add(
{"type": "response.created", "response": {"id": response_id}},
received_at_s=now + 0.01,
)
self.events.add(
{
"type": "response.audio_transcript.delta",
"response_id": response_id,
"delta": f"turn {self.response_count}",
},
received_at_s=now + 0.02,
)
self.events.add(
{
"type": "response.audio.delta",
"response_id": response_id,
"delta": base64.b64encode(b"\x00\x00" * 2400).decode(),
"sample_rate_hz": 24_000,
"metadata": {"audio_duration_ms": 100},
},
received_at_s=now + 0.03,
)
self.events.add(
{"type": "response.done", "response": {"id": response_id}},
received_at_s=now + 0.04,
)
async def acknowledge_playback(self):
self.ack_count += 1
async def close_session(self, **_kwargs):
return None
monkeypatch.setattr(
"vllm_omni.benchmarks.patch.patch.RealtimeDuplexClient",
FakeRealtimeClient,
)
request_input = RequestFuncInput(
model="openbmb/MiniCPM-o-4_5",
model_name="openbmb/MiniCPM-o-4_5",
prompt="hello",
api_url="http://localhost:8000/v1/realtime",
prompt_len=1,
output_len=20,
logprobs=None,
multi_modal_content=None,
ignore_eos=False,
extra_body={"save_duplex_request_metrics": True},
)
request_input.seed_tts_speech_extra = {"ref_audio": "data:audio/wav;base64,AAAA"}
request_input.seed_tts_system_prompt = "Speak exactly."
request_input.seed_tts_turns = tuple(
SimpleNamespace(utterance_id=f"utt-{index}", target_text=f"text {index}") for index in range(4)
)
output = await async_request_openai_realtime_duplex(
request_input,
session=None,
)
client = FakeRealtimeClient.last_instance
assert client.configure_kwargs["native_duplex"] is False
assert client.configure_kwargs["extra_body"] == {
"ref_audio": "data:audio/wav;base64,AAAA",
}
assert client.sent == [
event
for index in range(4)
for event in (
{
"type": "conversation.item.create",
"item": {
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": f"text {index}"}],
},
},
{"type": "response.create"},
)
]
assert client.ack_count == 4
assert output.success is True
assert output.generated_text == "turn 1 turn 2 turn 3 turn 4"
assert output.audio_duration == pytest.approx(0.4)
assert output.ttft > 0
assert output.audio_ttfp > output.ttft
assert output.audio_rtf > 0
assert output.latency > 0
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"]
assert len(output.duplex_request_metrics) == 4
for index, metrics in enumerate(output.duplex_request_metrics):
assert metrics == {
"session_id": session_id,
"request_index": index,
"utterance_id": f"utt-{index}",
"response_id": f"resp-{index + 1}",
"source": "client_monotonic_receive",
"measurement_origin": {
"ttft": "conversation.item.create client send to first non-empty text delta",
"ttfp": "conversation.item.create client send to first audio packet",
"rtf": "request-start-to-last-audio receive time divided by emitted audio duration",
},
"ttft_ms": pytest.approx(20.0, abs=2.0),
"ttfp_ms": pytest.approx(30.0, abs=2.0),
"rtf": pytest.approx(0.3, abs=0.03),
"audio_generation_ms": pytest.approx(30.0, abs=2.0),
"audio_duration_ms": 100.0,
}
assert output.duplex_session_metrics == {
"session_id": session_id,
"audio_turn_count": 4,
"mean_ttft_ms": pytest.approx(20.0, abs=2.0),
"mean_ttfp_ms": pytest.approx(30.0, abs=2.0),
"mean_rtf": pytest.approx(0.3, abs=0.03),
}
def create_sse_chunk(data_dict):
"""Helper to create SSE formatted chunk"""
return f"data: {json.dumps(data_dict)}\n\n".encode()
+1 -23
View File
@@ -42,10 +42,7 @@ from benchmarks.accuracy.text_to_image.gbench import (
summarize_generated_records as summarize_gebench_generated_records,
summarize_gebench_results,
)
from tests.e2e.accuracy.qwen3_omni.qwen3_omni_acc_bench_core import (
build_serve_common_argv,
seed_tts_bench_argv,
)
from tests.e2e.accuracy.qwen3_omni.qwen3_omni_acc_bench_core import seed_tts_bench_argv
from tests.e2e.accuracy.qwen3_omni.run_qwen_omni_acc_benchmark import sync_dataset_env_from_ns
from vllm_omni.benchmarks.data_modules.seed_tts_dataset import resolve_seed_tts_root
@@ -60,25 +57,6 @@ def test_seed_tts_bench_argv_preserves_hf_repo_id_from_env(monkeypatch):
assert argv[dataset_idx + 1] == "zhaochenyang20/seed-tts-eval"
def test_build_serve_common_argv_accepts_realtime_seed_tts_route(tmp_path):
argv = build_serve_common_argv(
host="127.0.0.1",
port=8000,
model="openbmb/MiniCPM-o-4_5",
num_prompts=4,
max_concurrency=1,
num_warmups=0,
percentile_metrics="ttft,audio_ttfp,audio_rtf",
result_dir=tmp_path,
result_filename="result.json",
backend="openai-realtime-tts",
endpoint="/v1/realtime",
)
assert argv[argv.index("--backend") + 1] == "openai-realtime-tts"
assert argv[argv.index("--endpoint") + 1] == "/v1/realtime"
def test_sync_dataset_env_preserves_seed_tts_hf_repo_id(monkeypatch):
ns = argparse.Namespace(
daily_omni_repo=None,
@@ -28,13 +28,10 @@ if _MODULE_NAME not in sys.modules:
_spec.loader.exec_module(_mod)
from vllm_omni.benchmarks.data_modules.seed_tts_dataset import ( # noqa: E402
SeedTTSDataset,
SeedTTSDesignDataset,
SeedTTSDesignSampleRequest,
SeedTTSSampleRequest,
SeedTTSTextDataset,
SeedTTSTextSampleRequest,
SeedTTSTurn,
)
# ---------------------------------------------------------------------------
@@ -73,94 +70,6 @@ def mock_tokenizer(mocker):
# ---------------------------------------------------------------------------
def test_seed_tts_dataset_groups_four_turns_with_one_reference(seed_tts_root, mock_tokenizer):
ds = SeedTTSDataset(
dataset_path=str(seed_tts_root),
random_seed=0,
locale="en",
disable_shuffle=True,
)
requests = ds.sample(
mock_tokenizer,
num_requests=1,
turns_per_session=4,
no_oversample=True,
)
assert len(requests) == 1
request = requests[0]
assert [turn.utterance_id for turn in request.seed_tts_turns] == [
"utt000",
"utt001",
"utt002",
"utt003",
]
assert [turn.target_text for turn in request.seed_tts_turns] == [
"target text 0",
"target text 1",
"target text 2",
"target text 3",
]
assert request.seed_tts_speech_extra["ref_text"] == "ref text 0"
assert request.seed_tts_ref_wav_path.endswith("utt000.wav")
def test_seed_tts_eval_expands_grouped_pcm_by_turn():
from vllm_omni.benchmarks.data_modules.seed_tts_eval import (
_expand_seed_tts_turn_outputs,
)
request = SeedTTSSampleRequest(
prompt="target 0",
prompt_len=4,
expected_output_len=100,
multi_modal_data=None,
request_id="session-0",
seed_tts_turns=tuple(
SeedTTSTurn(utterance_id=f"utt-{index}", target_text=f"target {index}") for index in range(4)
),
)
output = types.SimpleNamespace(
success=True,
tts_turn_pcm_bytes=[bytes([index]) for index in range(4)],
)
requests, outputs = _expand_seed_tts_turn_outputs([request], [output])
assert [item.prompt for item in requests] == [f"target {index}" for index in range(4)]
assert [item.seed_tts_utterance_id for item in requests] == [f"utt-{index}" for index in range(4)]
assert [item.tts_output_pcm_bytes for item in outputs] == [bytes([index]) for index in range(4)]
def test_seed_tts_eval_keeps_session_pcm_for_single_turn_chat():
"""Chat-omni Seed-TTS sets session PCM only; expand must not wipe it."""
from vllm_omni.benchmarks.data_modules.seed_tts_eval import (
_expand_seed_tts_turn_outputs,
)
request = SeedTTSSampleRequest(
prompt="target 0",
prompt_len=4,
expected_output_len=100,
multi_modal_data=None,
request_id="session-0",
seed_tts_turns=(SeedTTSTurn(utterance_id="utt-0", target_text="target 0"),),
)
pcm = b"\x01\x02\x03\x04"
output = types.SimpleNamespace(
success=True,
tts_output_pcm_bytes=pcm,
tts_turn_pcm_bytes=None,
)
requests, outputs = _expand_seed_tts_turn_outputs([request], [output])
assert len(requests) == 1
assert requests[0].prompt == "target 0"
assert outputs[0].tts_output_pcm_bytes == pcm
def test_seed_tts_text_dataset_omits_ref_audio(seed_tts_root, mock_tokenizer):
ds = SeedTTSTextDataset(
dataset_path=str(seed_tts_root),
+1 -2
View File
@@ -2241,8 +2241,7 @@ class TestPlatformOverrides:
load_deploy_config(Path(get_deploy_config_path(filename))), platform="cuda"
)
replica_stages = merge_pipeline_deploy(pipeline, replica)
# Explicit null clears the inherited single-GPU 2 GiB CUDA cap.
assert replica_stages[1].yaml_engine_args.get("kv_cache_memory_bytes") is None
assert replica_stages[1].yaml_engine_args["kv_cache_memory_bytes"] is None
def test_npu_overrides(self):
deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
-33
View File
@@ -39,7 +39,6 @@ from vllm_omni.config.stage_config import (
load_deploy_config,
merge_pipeline_deploy,
)
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -474,8 +473,6 @@ def test_sub_config_fields_match_structured_scopes():
"moe_backend",
"hf_overrides",
"limit_mm_per_prompt",
"interleave_mm_strings",
"media_io_kwargs",
"active_stream_window",
"duplex_max_sessions",
"enable_sleep_mode",
@@ -969,24 +966,20 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
diffusion_attention_backend="flash_attn",
fa_deterministic=True,
kv_cache_dtype="fp8",
kv_cache_skip_steps="0-1",
kv_cache_skip_layers=[2],
static_lora_scale=0.25,
diffusion_kv_mode="paged_scheduler",
diffusers_load_kwargs=None,
diffusers_call_kwargs=None,
)
assert cfg.diffusion_attention_config.default.backend == "flash_attn"
assert cfg.fa_deterministic is True
assert cfg.diffusion_kv_cache_dtype == "fp8"
assert cfg.diffusion_kv_cache_skip_step_indices == {0, 1}
assert cfg.diffusion_kv_cache_skip_layer_indices == {2}
assert cfg.lora_scale == 0.25
assert cfg.cache_backend == "tea_cache"
assert cfg.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
assert cfg.diffusers_load_kwargs == {}
assert cfg.diffusers_call_kwargs == {}
@@ -1001,8 +994,6 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
"stages:",
" - stage_id: 0",
" diffusion_attention_backend: flash_attn",
" fa_deterministic: true",
" diffusion_kv_mode: paged_scheduler",
]
)
)
@@ -1014,29 +1005,6 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
assert isinstance(stage, VllmOmniDiffusionStageConfig)
assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn"
assert stage.diffusion_config.fa_deterministic is True
assert stage.diffusion_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path):
deploy_path = tmp_path / "dreamzero_reserved_diffusion_kv.yaml"
deploy_path.write_text(
"\n".join(
[
"pipeline: dreamzero",
"async_chunk: false",
"stages:",
" - stage_id: 0",
" diffusion_kv_mode: paged_worker_local",
]
)
)
with pytest.raises(ValueError, match="reserved but not implemented"):
_from_pipeline_key(
"dreamzero",
deploy_config_path=str(deploy_path),
)
def test_diffusion_config_field_classification_covers_current_fields():
@@ -1057,7 +1025,6 @@ def test_diffusion_config_field_classification_covers_current_fields():
"enable_session_state_manager",
"prompt_embed_cache_size",
"diffusion_kv_cache_dtype",
"diffusion_kv_mode",
} <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
assert {
"revision",
-9
View File
@@ -752,12 +752,3 @@ def pytest_addoption(parser: pytest.Parser) -> None:
default=None,
help=("Path to benchmark config JSON. Example: --test-config-file tests/dfx/perf/tests/test_tts.json"),
)
parser.addoption(
"--assert-baseline",
action="store_true",
default=False,
help=(
"When set, omni/diffusion perf runners compare metrics against the baseline block in the JSON config "
"(default: off)."
),
)
+3 -15
View File
@@ -157,19 +157,8 @@ def benchmark_params(request):
}
def assert_result(result, params, num_prompt) -> None:
def assert_result(result, num_prompt) -> None:
assert result["completed"] == num_prompt, "Request failures exist"
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")
assert isinstance(session_metrics, list), "Duplex session metrics are missing"
assert len(session_metrics) == num_prompt, (
f"Expected {num_prompt} duplex session metric rows, got {len(session_metrics)}"
)
assert all(
isinstance(metric, dict) and metric.get("audio_turn_count") == expected_audio_turns
for metric in session_metrics
), f"Not every duplex session emitted {expected_audio_turns} audio turns"
@pytest.mark.benchmark
@@ -223,7 +212,6 @@ def test_performance_benchmark(omni_server, benchmark_params):
"enabled",
"eval_phase",
"trust_remote_code",
"expected_duplex_audio_turns_per_session",
}
for key, value in params.items():
@@ -263,7 +251,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
random_output_len=params.get("random_output_len"),
resource_label=resource_label,
)
assert_result(result, params, num_prompt)
assert_result(result, num_prompt)
# concurrency test
for sweep_index, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
@@ -280,4 +268,4 @@ def test_performance_benchmark(omni_server, benchmark_params):
random_output_len=params.get("random_output_len"),
resource_label=resource_label,
)
assert_result(result, params, num_prompt)
assert_result(result, num_prompt)
@@ -1,87 +0,0 @@
[
{
"test_name": "test_minicpmo_4_5",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100",
"npu": "A3"
},
"num_cards": 1
}
},
"full_model",
"omni"
],
"server_params": {
"model": "openbmb/MiniCPM-o-4_5",
"extra_cli_args": [
"--deploy-config",
"vllm_omni/deploy/minicpmo_4_5.yaml",
"--trust-remote-code"
]
},
"benchmark_params": [
{
"dataset_name": "seed-tts",
"dataset_path": "zhaochenyang20/seed-tts-eval",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [
32,
64,
128
],
"max_concurrency": [
1,
4,
8
],
"no_oversample": true,
"trust_remote_code": true,
"seed_tts_locale": "en",
"extra_body": {
"modalities": [
"text",
"audio"
],
"chat_template_kwargs": {
"enable_thinking": false,
"use_tts_template": true
}
},
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration",
"baseline": {
"H100": {
"request_throughput": [
0.7803,
1.1008,
1.2027
],
"mean_ttft_ms": [
166.9764,
307.7133,
323.584
],
"mean_e2el_ms": [
1361.1682,
3800.6332,
6961.8636
],
"mean_audio_ttfp_ms": [
635.8099,
1615.6591,
2888.9831
],
"mean_audio_rtf": [
0.3151,
0.8568,
1.5425
]
}
}
}
]
}
]
@@ -1,68 +0,0 @@
[
{
"test_name": "test_minicpmo_4_5_duplex_seed_tts",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100",
"npu": "A3"
},
"num_cards": 1
}
},
"full_model",
"omni"
],
"server_params": {
"model": "openbmb/MiniCPM-o-4_5",
"extra_cli_args": [
"--deploy-config",
"vllm_omni/deploy/minicpmo_4_5_duplex.yaml",
"--trust-remote-code"
]
},
"benchmark_params": [
{
"dataset_name": "seed-tts",
"dataset_path": "zhaochenyang20/seed-tts-eval",
"backend": "openai-realtime-tts",
"endpoint": "/v1/realtime",
"num_prompts": [
1
],
"max_concurrency": [
1
],
"no_oversample": true,
"trust_remote_code": true,
"seed_tts_locale": "en",
"seed_tts_turns_per_session": 4,
"expected_duplex_audio_turns_per_session": 4,
"extra_body": {
"save_duplex_request_metrics": true
},
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration",
"baseline": {
"H100": {
"request_throughput": [
0.1594
],
"mean_ttft_ms": [
274.164
],
"mean_e2el_ms": [
6250.7696
],
"mean_audio_ttfp_ms": [
994.422
],
"mean_audio_rtf": [
0.46163
]
}
}
}
]
}
]
@@ -485,30 +485,3 @@ def test_diffusion_build_run_params_resolves_baseline_per_sweep(tmp_path, monkey
assert sweeps[2]["params"]["baseline"]["H100"]["throughput_qps"] == 0.3
assert "A3" in sweeps[2]["params"]["baseline"]
assert not isinstance(sweeps[2]["params"]["baseline"]["H100"]["throughput_qps"], list)
def test_omni_duplex_expected_audio_turns_accepts_complete_session():
from tests.dfx.perf.scripts.run_benchmark import assert_result
assert_result(
{
"completed": 1,
"duplex_session_metrics": [{"audio_turn_count": 4}],
},
{"expected_duplex_audio_turns_per_session": 4},
1,
)
def test_omni_duplex_expected_audio_turns_rejects_incomplete_session():
from tests.dfx.perf.scripts.run_benchmark import assert_result
with pytest.raises(AssertionError, match="emitted 4 audio turns"):
assert_result(
{
"completed": 1,
"duplex_session_metrics": [{"audio_turn_count": 3}],
},
{"expected_duplex_audio_turns_per_session": 4},
1,
)
@@ -43,9 +43,9 @@
],
"baseline": {
"H100": {
"throughput_qps": 0.0361,
"latency_mean": 27.6206,
"peak_memory_mb_mean": 80548.0
"throughput_qps": 0.0434,
"latency_mean": 23.0654,
"peak_memory_mb_mean": 76360.0
}
}
}
@@ -100,9 +100,9 @@
],
"baseline": {
"H100": {
"throughput_qps": 0.05263,
"latency_mean": 18.9415,
"peak_memory_mb_mean": 50053.7
"throughput_qps": 0.0606,
"latency_mean": 16.4928,
"peak_memory_mb_mean": 47164.8
}
}
},
@@ -127,9 +127,9 @@
],
"baseline": {
"H100": {
"throughput_qps": 0.009595,
"latency_mean": 104.1058,
"peak_memory_mb_mean": 59718.0
"throughput_qps": 0.0124,
"latency_mean": 80.7784,
"peak_memory_mb_mean": 55374.3429
}
}
}
@@ -50,12 +50,12 @@ def _prompt() -> str:
return "What is the capital of China? Answer in 20 words."
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1)
@pytest.mark.parametrize("omni_server", _MINICPMO_SERVER, indirect=True)
def test_invalid_audio_format_rejected(omni_server: OmniServerParams, openai_client: OpenAIClientHandler) -> None:
"""
Test that invalid audio format is properly rejected.
Deploy Setting: default single GPU
Deploy Setting: default 2GPU
Input Modal: text + invalid audio format
Output Modal: error
"""
@@ -315,32 +315,11 @@ def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -
pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}")
def _assert_server_alive_after_fault(host: str, port: int, *, scenario: str) -> None:
"""#4285: a single-stage engine death must NOT take down the API server.
The process must stay up and keep answering /health (503 is acceptable —
errored but listening). A connection-refused means the server process
exited, which is the bug this fix closes.
"""
deadline = time.monotonic() + 15.0
responded = False
while time.monotonic() < deadline:
try:
status, _ = get_health_raw(host, port, timeout_sec=5)
responded = True
assert status in (200, 503), f"[{scenario} alive] unexpected /health status={status}"
except Exception as exc: # noqa: BLE001
if _looks_like_server_unreachable(exc):
pytest.fail(f"[{scenario} alive] server exited after single-stage OOM (#4285): {exc!r}")
raise
time.sleep(0.5)
assert responded, f"[{scenario} alive] /health never responded after fault"
QWEN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR)
@pytest.mark.slow
@pytest.mark.skip(reason="issue#4285")
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
@@ -407,11 +386,9 @@ def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
assert "code" in chat_error, f"chat error lacks code field: {chat_error!r}"
assert "code" in speech_error, f"speech error lacks code field: {speech_error!r}"
# #4285: a single dead stage must not take the whole API server down.
_assert_server_alive_after_fault(host, port, scenario="oom_contract")
@pytest.mark.slow
@pytest.mark.skip(reason="issue#4285")
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, openai_client_function) -> None:
@@ -454,6 +431,7 @@ def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_functi
@pytest.mark.slow
@pytest.mark.skip(reason="issue#4285")
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, openai_client_function) -> None:
@@ -796,11 +774,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
@pytest.mark.slow
@pytest.mark.skip(
reason="needs request-scoped OOM recovery: under transient pressure the stage "
"EngineCore still dies and is evicted (not respawned), so the engine does not "
"recover after the hog stops. Out of scope for the #4285 fault-isolation fix."
)
@pytest.mark.skip(reason="issue#4285")
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
def test_reliability_fault_gpu_oom_state_converges_after_fault_removed(
@@ -9,10 +9,6 @@ This script tests two main scenarios:
2. Case 2: Comparing FlashAttention and SDPA backends for batch_size=2 with padding
"""
import sys
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
@@ -375,264 +371,6 @@ def test_packed_varlen_metadata_must_be_complete(monkeypatch):
impl.forward_cuda(query, query, query, metadata)
# ---------------------------------------------------------------------------
# NPU mask-free packed attention (PR #5891).
#
# These tests cover the backend-side NPU branches that ``forward_fa_npu`` adds:
# - boundary resolution in ``_resolve_packed_seq_npu`` (pure Python, no device
# sync, no MindIE-SD import);
# - env dispatch in ``forward_fa_npu`` (which op a packed forward takes);
# - the laser input pre-scaling math in ``_forward_prefix_kv_slice_npu``.
#
# They run on CPU without a real NPU by injecting a fake ``mindiesd`` module
# into ``sys.modules`` (same pattern as ``test_paged_attention`` /
# ``benchmarks/conftest``). No production code is exercised for real kernels.
def _npu_impl(causal: bool = False) -> FlashAttentionImpl:
return FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=causal)
def _fake_mindiesd(monkeypatch, *, attention_forward=None, attention_forward_varlen=None):
"""Install a stub ``mindiesd`` so local ``from mindiesd import ...`` works.
Restored automatically by monkeypatch. Defaults are no-op callables so the
top-level import in ``forward_fa_npu`` succeeds even when the test mocks the
instance methods that would otherwise call them.
"""
fake = SimpleNamespace(
attention_forward=attention_forward or Mock(return_value=torch.zeros(1)),
attention_forward_varlen=attention_forward_varlen or Mock(return_value=torch.zeros(1)),
)
monkeypatch.setitem(sys.modules, "mindiesd", fake)
return fake
# --- Test group A: boundary resolution (_resolve_packed_seq_npu) -------------
def test_resolve_packed_seq_accepts_real_plus_pad_contract():
impl = _npu_impl()
q = torch.randn(1, 8, 2, 4)
cu = torch.tensor([0, 5, 8], dtype=torch.int32)
extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 5, "max_seqlen_k": 5}
assert impl._resolve_packed_seq_npu(q, q, extra) == ([5, 8], [5, 8])
def test_resolve_packed_seq_single_document_no_padding():
impl = _npu_impl()
q = torch.randn(1, 8, 2, 4)
cu = torch.tensor([0, 8], dtype=torch.int32)
extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 8, "max_seqlen_k": 8}
assert impl._resolve_packed_seq_npu(q, q, extra) == ([8], [8])
@pytest.mark.parametrize(
"case",
[
"causal",
"missing_key",
"batch_ne_1",
"cu_too_long",
"cu_qk_mismatch",
"max_seqlen_not_int",
"used_gt_total",
"used_zero",
"pad_longer_than_real",
],
)
def test_resolve_packed_seq_rejects_malformed_metadata(case):
"""Any metadata outside the exact [real, pad] contract returns None so the
caller falls back to the masked path."""
cu3 = torch.tensor([0, 5, 8], dtype=torch.int32)
cu2 = torch.tensor([0, 8], dtype=torch.int32)
cu4 = torch.tensor([0, 3, 5, 8], dtype=torch.int32)
cu_short = torch.tensor([0, 3, 8], dtype=torch.int32)
base = {"cu_seqlens_q": cu3, "cu_seqlens_k": cu3, "max_seqlen_q": 5, "max_seqlen_k": 5}
q = torch.randn(1, 8, 2, 4)
impl = _npu_impl(causal=(case == "causal"))
if case == "missing_key":
extra = {"cu_seqlens_q": cu3, "max_seqlen_q": 5} # 3 of 4 keys absent
elif case == "batch_ne_1":
q = torch.randn(2, 8, 2, 4)
extra = base
elif case == "cu_too_long":
extra = {**base, "cu_seqlens_q": cu4, "cu_seqlens_k": cu4}
elif case == "cu_qk_mismatch":
extra = {**base, "cu_seqlens_q": cu3, "cu_seqlens_k": cu2}
elif case == "max_seqlen_not_int":
extra = {**base, "max_seqlen_q": 5.0}
elif case == "used_gt_total":
extra = {**base, "max_seqlen_q": 9} # 9 > total length 8
elif case == "used_zero":
extra = {**base, "max_seqlen_q": 0, "max_seqlen_k": 0}
elif case == "pad_longer_than_real":
# real=3, pad=5: the real document must be the longer one.
extra = {
"cu_seqlens_q": cu_short,
"cu_seqlens_k": cu_short,
"max_seqlen_q": 3,
"max_seqlen_k": 3,
}
else:
extra = base
assert impl._resolve_packed_seq_npu(q, q, extra) is None
# --- Test group B: env dispatch in forward_fa_npu (current behavior) --------
def _npu_impl_with_mocked_paths() -> tuple[FlashAttentionImpl, torch.Tensor]:
"""FlashAttentionImpl whose two packed paths are Mocks returning a sentinel.
Lets a test assert which branch ``forward_fa_npu`` dispatched to without
running any real kernel.
"""
impl = _npu_impl()
sentinel = torch.tensor([1.0])
impl._forward_varlen_packed_npu = Mock(return_value=sentinel)
impl._forward_prefix_kv_slice_npu = Mock(return_value=sentinel)
return impl, sentinel
def test_npu_env_unset_uses_varlen(monkeypatch):
monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False)
_fake_mindiesd(monkeypatch)
impl, _ = _npu_impl_with_mocked_paths()
metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
impl._forward_varlen_packed_npu.assert_called_once()
impl._forward_prefix_kv_slice_npu.assert_not_called()
def test_npu_env_laser_uses_slice(monkeypatch):
monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
_fake_mindiesd(monkeypatch)
impl, _ = _npu_impl_with_mocked_paths()
metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
impl._forward_prefix_kv_slice_npu.assert_called_once()
impl._forward_varlen_packed_npu.assert_not_called()
@pytest.mark.parametrize("fa_type", ["prompt_flash_attn", "fused_attn_score", "ascend_laser_attenton", "garbage"])
def test_npu_env_non_laser_value_falls_back_to_varlen(monkeypatch, fa_type):
"""Locks in the CURRENT dispatch behavior (PR #5891).
Any value other than ``ascend_laser_attention`` — including the valid
``prompt_flash_attn``/``fused_attn_score`` and typos — is silently routed to
the TND varlen op, which cannot honor ``MINDIE_SD_FA_TYPE``. This is review
comment #1 on the PR and is intentionally captured here as a regression
net: update this test when the dispatch starts honoring/rejecting these.
"""
monkeypatch.setenv("MINDIE_SD_FA_TYPE", fa_type)
_fake_mindiesd(monkeypatch)
impl, _ = _npu_impl_with_mocked_paths()
metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata)
impl._forward_varlen_packed_npu.assert_called_once()
impl._forward_prefix_kv_slice_npu.assert_not_called()
def test_npu_varlen_opt_in_unset_takes_mask_path(monkeypatch):
"""Models that do not set ``npu_attn_varlen`` (e.g. Wan/Cosmos) keep using
the existing masked attention_forward path; the new branches are skipped."""
fake_forward = Mock(return_value=torch.zeros(1, 8, 2, 4))
_fake_mindiesd(monkeypatch, attention_forward=fake_forward)
impl, _ = _npu_impl_with_mocked_paths()
metadata = AttentionMetadata(
attn_mask=torch.tensor([[True, True, True, True, True, False, False, False]]),
extra={}, # no npu_attn_varlen opt-in
)
q = torch.randn(1, 8, 2, 4)
out = impl.forward_fa_npu(q, q, q, metadata)
impl._forward_varlen_packed_npu.assert_not_called()
impl._forward_prefix_kv_slice_npu.assert_not_called()
fake_forward.assert_called_once()
assert out is fake_forward.return_value
# --- Test group C: laser input pre-scaling in _forward_prefix_kv_slice_npu --
def test_prefix_kv_slice_applies_laser_input_scaling(monkeypatch):
monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
captured: dict = {}
def fake_attention_forward(query, key, value, **kwargs):
captured["query"] = query
captured["key"] = key
captured["value"] = value
captured.update(kwargs)
return torch.ones_like(query) # deterministic stand-in kernel output
_fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
impl = _npu_impl() # softmax_scale == 0.5
q = torch.randn(1, 8, 2, 4)
cu = torch.tensor([0, 5, 8], dtype=torch.int32)
extra = {
"cu_seqlens_q": cu,
"cu_seqlens_k": cu,
"max_seqlen_q": 5,
"max_seqlen_k": 5,
"laser_input_scale": 256.0,
}
out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
# K/V sliced to the valid prefix (used_k == 5); q keeps full length.
assert captured["key"].shape == (1, 5, 2, 4)
assert captured["value"].shape == (1, 5, 2, 4)
assert captured["query"].shape == (1, 8, 2, 4)
assert captured["attn_mask"] is None
# Pre-divide q/k/v by the factor and compensate the kernel scale by its
# square (exact for power-of-two); output scaled back by the factor.
torch.testing.assert_close(captured["query"], q / 256.0)
torch.testing.assert_close(captured["key"], q[:, :5] / 256.0)
assert captured["scale"] == pytest.approx(0.5 * 256.0**2)
torch.testing.assert_close(out, torch.ones(1, 8, 2, 4) * 256.0)
def test_prefix_kv_slice_no_scaling_without_factor(monkeypatch):
monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
captured: dict = {}
def fake_attention_forward(query, key, value, **kwargs):
captured["query"] = query
captured.update(kwargs)
return torch.ones_like(query)
_fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward)
impl = _npu_impl()
q = torch.randn(1, 8, 2, 4)
cu = torch.tensor([0, 5, 8], dtype=torch.int32)
extra = {
"cu_seqlens_q": cu,
"cu_seqlens_k": cu,
"max_seqlen_q": 5,
"max_seqlen_k": 5,
# no laser_input_scale → no pre-scaling
}
out = impl._forward_prefix_kv_slice_npu(q, q, q, extra)
torch.testing.assert_close(captured["query"], q)
assert captured["scale"] == pytest.approx(0.5)
torch.testing.assert_close(out, torch.ones(1, 8, 2, 4))
if __name__ == "__main__":
print("Running FlashAttention Padding Tests...")
print("=" * 60)
@@ -5,9 +5,10 @@
The grids here are the ones MiniMax-H3 FL2VA actually produces for an 8.7s clip:
1280x768 gives a 62x24x40 latent grid whose 59520 video rows land on the 128-row
kernel block, and 1344x768 gives 62x24x42 whose 62496 rows do not. Both are
handed to MindIE-SD: its rf_v2 preprocessing promotes an irregular video tail
to the always-kept prefix segment before generating the block mask.
kernel block, and 1344x768 gives 62x24x42 whose 62496 rows straddle it. Only the
first can run sparse: rf_v2 tiles the sequence as a whole but pools its block
mask from the video and prefix segments separately, so the two only agree on the
block count when the video segment ends on a boundary.
"""
import dataclasses
@@ -56,13 +57,10 @@ def test_block_aligned_video_segment_runs_sparse():
assert plan.latent_shape == list(ALIGNED_GRID)
def test_misaligned_video_segment_runs_sparse_with_updated_mindiesd():
plan = make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID))
assert plan is not None
assert plan.prefix_len == PREFIX_ROWS
assert plan.used_len == PREFIX_ROWS + 62496
assert plan.latent_shape == list(MISALIGNED_GRID)
def test_misaligned_video_segment_stays_dense():
# Realigning the sequence would mean padding it, and rf_v2 ignores attn_mask:
# pad keys would take a share of every real query's softmax denominator.
assert make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is None
@pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900])
@@ -157,28 +155,3 @@ def test_fully_populated_mask_reproduces_dense_attention():
).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"
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="rf_v2 runs on Ascend NPU only.")
def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd():
"""vLLM forwards the full irregular grid; MindIE-SD owns its tail handling."""
torch.manual_seed(1)
prefix_len, grid = 384, (4, 16, 10)
video_len = grid[0] * grid[1] * grid[2]
used = prefix_len + video_len
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()
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)
reference = torch.nn.functional.scaled_dot_product_attention(
q.transpose(1, 2),
k.transpose(1, 2),
v.transpose(1, 2),
scale=impl.softmax_scale,
).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"
-60
View File
@@ -471,66 +471,6 @@ class TestTeaCacheBackend:
assert backend.enabled is True
mock_apply_hook.assert_called_once()
@patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
def test_enable_uses_generic_default_threshold(self, mock_apply_hook):
pipeline = Mock()
pipeline.__class__.__name__ = "QwenImagePipeline"
pipeline.transformer = Mock()
pipeline.transformer.__class__.__name__ = "QwenImageTransformer2DModel"
TeaCacheBackend(DiffusionCacheConfig()).enable(pipeline)
assert mock_apply_hook.call_args.args[1].rel_l1_thresh == 0.2
@pytest.mark.parametrize("partition", ["fl2va", "combined"])
@pytest.mark.parametrize(
("configured_threshold", "expected_threshold"),
[(None, 0.17), (0.2, 0.2)],
)
@patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
def test_minimax_h3_only_enables_fl2va_teacache(
self,
mock_apply_hook,
partition,
configured_threshold,
expected_threshold,
):
pipeline = Mock()
pipeline.__class__.__name__ = "MiniMaxH3Pipeline"
pipeline.partition = partition
pipeline.transformer = Mock()
pipeline.transformer.__class__.__name__ = "MiniMaxH3DiTModel"
pipeline.transformers_ref = Mock()
config = (
DiffusionCacheConfig()
if configured_threshold is None
else DiffusionCacheConfig(rel_l1_thresh=configured_threshold)
)
backend = TeaCacheBackend(config)
backend.enable(pipeline)
mock_apply_hook.assert_called_once()
transformer, teacache_config = mock_apply_hook.call_args.args
assert transformer is pipeline.transformer
assert transformer is not pipeline.transformers_ref
assert teacache_config.rel_l1_thresh == expected_threshold
@patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
def test_minimax_h3_rejects_ref2va_teacache(self, mock_apply_hook):
pipeline = Mock()
pipeline.__class__.__name__ = "MiniMaxH3Pipeline"
pipeline.partition = "ref2va"
pipeline.transformer = Mock()
pipeline.transformer.__class__.__name__ = "MiniMaxH3DiTModel"
backend = TeaCacheBackend(DiffusionCacheConfig())
with pytest.raises(ValueError, match="only supports the MiniMax-H3 FL2VA partition"):
backend.enable(pipeline)
assert backend.enabled is False
mock_apply_hook.assert_not_called()
@patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
def test_enable_with_coefficients(self, mock_apply_hook):
"""Test enabling TeaCache with custom coefficients."""
-289
View File
@@ -15,23 +15,19 @@ Currently implemented:
- TestFlux2KleinExtractor: Flux2Klein model extractor
- TestFlux2Extractor: Flux2 model extractor
- TestFluxExtractor: Flux model extractor
- TestMiniMaxH3Extractor: MiniMaxH3DiTModel extractor
"""
from abc import ABC, abstractmethod
from types import SimpleNamespace
from unittest.mock import MagicMock, Mock, patch
import pytest
import torch
import torch.nn as nn
from tests.helpers.mark import hardware_test
from vllm_omni.diffusion.cache.teacache.extractors import (
extract_flux2_context,
extract_flux2_klein_context,
extract_flux_context,
extract_minimax_h3_context,
)
from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformer2DModel
from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
@@ -427,288 +423,3 @@ class TestFluxExtractor(BaseExtractorTest):
img_ids=torch.randint(0, 64, (1, 16, 3)),
txt_ids=torch.randint(0, 64, (1, 8, 3)),
)
class _MiniMaxH3FakeLinear(nn.Module):
def __init__(
self,
in_features,
out_features,
*,
bias=False,
params_dtype=None,
quant_config=None,
prefix="",
total_num_heads=None,
total_num_kv_heads=None,
**kwargs,
):
del kwargs
super().__init__()
self.prefix = prefix
self.quant_config = quant_config
self.weight = nn.Parameter(torch.randn(out_features, in_features, dtype=params_dtype or torch.float32))
if bias:
self.bias = nn.Parameter(torch.zeros(out_features, dtype=params_dtype or torch.float32))
else:
self.register_parameter("bias", None)
self.num_heads = total_num_heads
self.num_kv_heads = total_num_kv_heads
def forward(self, x):
return torch.nn.functional.linear(x, self.weight, self.bias), None
class _MiniMaxH3FakeMergedLinear(_MiniMaxH3FakeLinear):
def __init__(self, in_features, out_features, **kwargs):
super().__init__(in_features, sum(out_features), **kwargs)
class _MiniMaxH3FakeQKVLinear(_MiniMaxH3FakeLinear):
def __init__(
self,
*,
hidden_size,
head_size,
total_num_heads,
total_num_kv_heads,
**kwargs,
):
qkv_features = (total_num_heads + 2 * total_num_kv_heads) * head_size
super().__init__(
hidden_size,
qkv_features,
total_num_heads=total_num_heads,
total_num_kv_heads=total_num_kv_heads,
**kwargs,
)
class _MiniMaxH3FakeAttention(nn.Module):
class FakeBackend:
supports_prefix_kv_slicing = False
# Mirror AttentionBackend.supports_packed_mask_free: H3's packed
# attention path calls attn_backend.supports_packed_mask_free() to decide
# whether to build the padding mask. Return False so the fake exercises
# the masked path (the abstract backend default).
@classmethod
def supports_packed_mask_free(cls) -> bool:
return False
def __init__(self, **kwargs):
del kwargs
super().__init__()
self.attn_backend = self.FakeBackend
self.use_ring = False
def forward(self, q, k, v, metadata=None):
del k, v, metadata
return q
def _minimax_h3_small_od_config():
arch = {
"num_layers": 1,
"token_refiner_num_layers": 1,
"hidden_size": 8,
"num_attention_heads": 1,
"attention_head_dim": 8,
"ffn_hidden_size": 16,
"latents_dim": 2,
"audio_latents_dim": 2,
"patch_size": (1, 2, 2),
"text_dim": 6,
"timestep_input_dim": 4,
"time_embed_hidden_size": 8,
"time_embed_dim": 4,
"adaln_out_features": 18 * 8,
"final_adaln_out_features": 2 * 8,
"rope_inv_freq_len": 1,
}
return SimpleNamespace(
tf_model_config=arch,
parallel_config=SimpleNamespace(ulysses_degree=1),
)
def _minimax_h3_sample_kwargs(*, seq_len: int = 8, device: torch.device | str = "cpu") -> dict:
video_row_dim = 2 * 1 * 2 * 2
audio_row_dim = 2
text_len = 2
img_rows = 2
img_pos = torch.tensor([2, 3], dtype=torch.long, device=device)
audio_pos = torch.tensor([4, 5], dtype=torch.long, device=device)
text_pos = torch.tensor([0, 1], dtype=torch.long, device=device)
infer_out_pos = img_pos.clone()
token_tags = torch.full((seq_len,), -1, dtype=torch.long, device=device)
token_tags[text_pos] = 1
token_tags[img_pos] = 0
token_tags[audio_pos] = 2
inverse_indices = torch.zeros(seq_len, dtype=torch.long, device=device)
return {
"x": torch.randn(1, seq_len, video_row_dim, device=device),
"audio_x": torch.randn(1, seq_len, audio_row_dim, device=device),
"img_position_ids": torch.zeros(1, seq_len, 3, dtype=torch.float64, device=device),
"unique_timesteps": torch.tensor([0.5], dtype=torch.float32, device=device),
"inverse_indices": inverse_indices,
"update_mask": torch.ones(img_rows, dtype=torch.bool, device=device),
"token_tags": token_tags,
"prompt_embeds": torch.randn(text_len, 6, device=device),
"img_pos_info": {"position_ids": img_pos},
"audio_pos_info": {"position_ids": audio_pos},
"text_pos_info": {"position_ids": text_pos},
"img_pos_for_infer_output_info": {"position_ids": infer_out_pos},
"packed_seq_params": {
"cu_seqlens_q": torch.tensor([0, seq_len], dtype=torch.int32, device=device),
"max_seqlen_q": seq_len,
},
"refiner_packed_seq_params": {
"cu_seqlens_q": torch.tensor([0, text_len], dtype=torch.int32, device=device),
"max_seqlen_q": text_len,
},
}
@pytest.mark.cpu
class TestMiniMaxH3Extractor(BaseExtractorTest):
"""Test extract_minimax_h3_context function."""
def get_extractor(self):
return extract_minimax_h3_context
@pytest.fixture
def minimax_h3_module(self, monkeypatch):
"""Create a minimal MiniMaxH3DiTModel with fake parallel linears for CPU tests."""
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
monkeypatch.setattr(h3, "ColumnParallelLinear", _MiniMaxH3FakeLinear)
monkeypatch.setattr(h3, "MergedColumnParallelLinear", _MiniMaxH3FakeMergedLinear)
monkeypatch.setattr(h3, "QKVParallelLinear", _MiniMaxH3FakeQKVLinear)
monkeypatch.setattr(h3, "RowParallelLinear", _MiniMaxH3FakeLinear)
monkeypatch.setattr(h3, "Attention", _MiniMaxH3FakeAttention)
monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1)
model = h3.MiniMaxH3DiTModel(_minimax_h3_small_od_config(), quant_config=None)
for submodule in model.modules():
if isinstance(submodule, h3.MiniMaxH3Attention):
submodule.rope._forward_method = submodule.rope.forward_native
model.eval()
return model
def get_module(self, minimax_h3_module):
return minimax_h3_module
@pytest.fixture
def sample_inputs(self):
"""Create sample packed forward kwargs for MiniMaxH3DiTModel."""
return _minimax_h3_sample_kwargs(seq_len=8)
def get_sample_inputs(self, sample_inputs):
return sample_inputs
def test_modulated_input_shape(self, minimax_h3_module, sample_inputs):
"""Test that modulated_input matches packed sequence shape [T, H]."""
context = extract_minimax_h3_context(minimax_h3_module, **sample_inputs)
seq_len = sample_inputs["x"].shape[1]
assert context.modulated_input.shape == (seq_len, minimax_h3_module.hidden_size)
def test_run_transformer_blocks_callable(self, minimax_h3_module, sample_inputs):
"""Test that run_transformer_blocks is callable."""
context = extract_minimax_h3_context(minimax_h3_module, **sample_inputs)
assert callable(context.run_transformer_blocks)
def test_postprocess_callable(self, minimax_h3_module, sample_inputs):
"""Test that postprocess is callable."""
context = extract_minimax_h3_context(minimax_h3_module, **sample_inputs)
assert callable(context.postprocess)
def test_postprocess_output_shape(self, minimax_h3_module, sample_inputs):
"""Test that postprocess returns video/audio logits for infer rows."""
context = extract_minimax_h3_context(minimax_h3_module, **sample_inputs)
(hidden,) = context.run_transformer_blocks()
video_logits, audio_logits = context.postprocess(hidden)
num_video_rows = sample_inputs["img_pos_info"]["position_ids"].shape[0]
num_audio_rows = sample_inputs["audio_pos_info"]["position_ids"].shape[0]
assert video_logits.shape[0] == num_video_rows
assert audio_logits.shape[0] == num_audio_rows
def test_extractor_matches_native_forward(self, minimax_h3_module, sample_inputs):
"""The uncached extractor path must preserve the model's forward result."""
expected_video, expected_audio = minimax_h3_module(**sample_inputs)
context = extract_minimax_h3_context(minimax_h3_module, **sample_inputs)
(hidden,) = context.run_transformer_blocks()
actual_video, actual_audio = context.postprocess(hidden)
torch.testing.assert_close(actual_video, expected_video)
torch.testing.assert_close(actual_audio, expected_audio)
def test_run_transformer_blocks_forwards_packed_layout(
self,
minimax_h3_module,
sample_inputs,
monkeypatch,
):
"""The extractor block call must mirror the native packed contract."""
from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout
video_layout = VideoTokenLayout(prefix_len=2, latent_grid=(1, 1, 2))
inputs = {**sample_inputs, "video_token_layout": video_layout}
block = minimax_h3_module.blocks[0]
original_forward = block.forward
received: list[tuple[int, VideoTokenLayout | None]] = []
def capture_forward(*args, **kwargs):
received.append((kwargs["packed_total"], kwargs["video_layout"]))
return original_forward(*args, **kwargs)
monkeypatch.setattr(block, "forward", capture_forward)
context = extract_minimax_h3_context(minimax_h3_module, **inputs)
context.run_transformer_blocks()
assert received == [(sample_inputs["x"].shape[1], video_layout)]
def test_teacache_reuses_h3_block_residual(self, minimax_h3_module, sample_inputs, monkeypatch):
"""A guaranteed cache hit must skip the H3 transformer block."""
from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook
block = minimax_h3_module.blocks[0]
original_forward = block.forward
block_calls = 0
def counted_forward(*args, **kwargs):
nonlocal block_calls
block_calls += 1
return original_forward(*args, **kwargs)
monkeypatch.setattr(block, "forward", counted_forward)
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)
video_logits, audio_logits = hook.new_forward(minimax_h3_module, **sample_inputs)
assert block_calls == 1
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]
def test_invalid_module_raises_error(self):
"""Test that invalid module without blocks raises ValueError."""
invalid_module = Mock()
invalid_module.blocks = []
with pytest.raises(ValueError, match="Module must have blocks"):
extract_minimax_h3_context(invalid_module, **_minimax_h3_sample_kwargs())
@@ -1,52 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@pytest.fixture(autouse=True)
def _fixed_master_port(monkeypatch) -> None:
monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
def test_dense_legacy_is_default() -> None:
config = OmniDiffusionConfig.from_kwargs()
assert config.diffusion_kv_mode is DiffusionKVCacheMode.DENSE_LEGACY
def test_paged_worker_local_is_rejected_until_implemented() -> None:
with pytest.raises(ValueError, match="reserved but not implemented"):
OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="paged_worker_local")
def test_unknown_cache_mode_is_rejected() -> None:
with pytest.raises(ValueError, match="Unsupported Diffusion KV diffusion_kv_mode"):
OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="unknown")
def test_non_mapping_omni_kv_config_is_rejected() -> None:
with pytest.raises(TypeError, match="omni_kv_config must be a mapping"):
OmniDiffusionConfig.from_kwargs(omni_kv_config="paged_scheduler")
def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None:
with pytest.raises(ValueError, match="does not support imported AR KV"):
OmniDiffusionConfig.from_kwargs(
diffusion_kv_mode="paged_scheduler",
omni_kv_config={"need_recv_cache": True},
)
def test_paged_scheduler_does_not_depend_on_model_registry() -> None:
config = OmniDiffusionConfig.from_kwargs(
model_class_name="FutureDiffusionModel",
diffusion_kv_mode="paged_scheduler",
)
assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
@@ -1,110 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm_omni.diffusion.diffusion_kv.metadata import (
DiffusionKVContextMetadata,
DiffusionKVMetadata,
DiffusionKVSequenceMetadata,
)
pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
def test_diffusion_kv_metadata_uses_native_cache_group_block_ids() -> None:
context = DiffusionKVContextMetadata(
context_id="text",
cache_role="cross_attention",
num_tokens=3,
block_ids=([7], [11]),
)
sequence = DiffusionKVSequenceMetadata(
sequence_id=1,
prefix_len=4,
target_len=2,
seq_len=8,
block_ids=([1, 2], [5, 6]),
context_ids=(context.context_id,),
)
metadata = DiffusionKVMetadata(
request_id="req-0",
allocation_generation=3,
sequences=(sequence,),
contexts=(context,),
)
assert metadata.sequences[0].block_ids == ([1, 2], [5, 6])
assert metadata.contexts[0].block_ids == ([7], [11])
def test_cfg_sequences_can_share_one_request_context() -> None:
shared_context = DiffusionKVContextMetadata(
context_id="shared-text",
cache_role="cross_attention",
num_tokens=3,
block_ids=([7], [11]),
)
sequences = tuple(
DiffusionKVSequenceMetadata(
sequence_id=sequence_id,
prefix_len=4,
target_len=2,
seq_len=8,
block_ids=([sequence_id + 1],),
context_ids=(shared_context.context_id,),
)
for sequence_id in range(2)
)
metadata = DiffusionKVMetadata(
request_id="req-shared-context",
allocation_generation=1,
sequences=sequences,
contexts=(shared_context,),
)
assert metadata.contexts == (shared_context,)
assert metadata.sequences[0].context_ids == (shared_context.context_id,)
assert metadata.sequences[1].context_ids == (shared_context.context_id,)
def test_cfg_sequences_can_reference_branch_specific_contexts() -> None:
conditional_context = DiffusionKVContextMetadata(
context_id="conditional-text",
cache_role="cross_attention",
num_tokens=3,
block_ids=([7], [11]),
)
unconditional_context = DiffusionKVContextMetadata(
context_id="unconditional-text",
cache_role="cross_attention",
num_tokens=3,
block_ids=([8], [12]),
)
metadata = DiffusionKVMetadata(
request_id="req-branch-contexts",
allocation_generation=1,
sequences=(
DiffusionKVSequenceMetadata(
sequence_id=0,
prefix_len=4,
target_len=2,
seq_len=8,
block_ids=([1],),
context_ids=(conditional_context.context_id,),
),
DiffusionKVSequenceMetadata(
sequence_id=1,
prefix_len=4,
target_len=2,
seq_len=8,
block_ids=([2],),
context_ids=(unconditional_context.context_id,),
),
),
contexts=(conditional_context, unconditional_context),
)
assert metadata.sequences[0].context_ids == (conditional_context.context_id,)
assert metadata.sequences[1].context_ids == (unconditional_context.context_id,)
@@ -1,200 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
from vllm.v1.core.kv_cache_manager import KVCacheManager
from vllm.v1.core.kv_cache_utils import BlockHash
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
KVCacheTensor,
)
from vllm.v1.request import RequestStatus
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
BLOCK_SIZE = 4
def _request(**overrides) -> DiffusionKVRequest:
values = dict(
request_id="req/diffusion-kv/0",
sequence_id=0,
prefix_len=4,
target_len=4,
seq_len=8,
)
values.update(overrides)
return DiffusionKVRequest(**values)
def test_request_exposes_native_request_and_diffusion_semantics() -> None:
request = _request()
assert request.sequence_id == 0
assert request.prefix_len == 4
assert request.target_len == 4
assert request.num_tokens == request.seq_len == 8
assert request.num_prompt_tokens == 4
assert request.num_computed_tokens == 0
assert request.block_hashes == []
assert request.kv_contexts == ()
assert request.skip_reading_prefix_cache is True
assert request.shared_prefix_boundary == 0
assert request.status is RequestStatus.WAITING
@pytest.mark.parametrize(
("overrides", "message"),
[
({"sequence_id": -1}, "sequence_id"),
({"prefix_len": -1}, "prefix_len"),
({"target_len": 0}, "target_len"),
({"seq_len": 0}, "seq_len"),
({"prefix_len": 5, "target_len": 4, "seq_len": 8}, "must not exceed"),
],
)
def test_request_rejects_invalid_lengths(overrides, message: str) -> None:
with pytest.raises(ValueError, match=message):
_request(**overrides)
@pytest.mark.parametrize(
("overrides", "message"),
[
({"context_id": ""}, "context_id"),
({"cache_role": ""}, "cache_role"),
({"num_tokens": 0}, "num_tokens"),
],
)
def test_context_rejects_invalid_identity_or_length(overrides, message: str) -> None:
values = dict(context_id="text", cache_role="cross.text", num_tokens=8)
values.update(overrides)
with pytest.raises(ValueError, match=message):
DiffusionKVContext(**values)
def test_request_owns_independent_kv_contexts() -> None:
text = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
image = DiffusionKVContext(context_id="image", cache_role="cross.image", num_tokens=4)
request = _request(kv_contexts=(text, image))
assert request.kv_contexts == (text, image)
def test_request_rejects_duplicate_context_ids() -> None:
contexts = (
DiffusionKVContext(context_id="condition", cache_role="cross.text", num_tokens=8),
DiffusionKVContext(context_id="condition", cache_role="cross.image", num_tokens=4),
)
with pytest.raises(ValueError, match="unique context_id"):
_request(kv_contexts=contexts)
def test_request_rejects_non_context_values() -> None:
with pytest.raises(TypeError, match="DiffusionKVContext"):
_request(kv_contexts=(object(),))
def _manager(*, enable_caching: bool) -> KVCacheManager:
# vLLM 0.26 lazily registers its built-in KV specs. Another test may have
# already registered an out-of-tree spec, making the registry non-empty
# before FullAttentionSpec is registered. Initialize the native mappings
# explicitly; vLLM 0.25 has no registration hook and needs no setup.
register_all = getattr(native_kv_managers, "register_all_kvcache_specs", None)
if register_all is not None:
register_all(None)
spec = FullAttentionSpec(
block_size=BLOCK_SIZE,
num_kv_heads=2,
head_size=8,
dtype=torch.bfloat16,
)
group = KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)
config = KVCacheConfig(
num_blocks=8,
kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["layer0"])],
kv_cache_groups=[group],
)
return KVCacheManager(
config,
max_model_len=64,
scheduler_block_size=BLOCK_SIZE,
hash_block_size=BLOCK_SIZE,
enable_caching=enable_caching,
)
def test_empty_hash_request_conforms_to_native_request_local_allocation() -> None:
manager = _manager(enable_caching=True)
free_before = manager.block_pool.get_num_free_blocks()
request = _request()
# vLLM 0.26 adds a shared-prefix-boundary return value.
computed_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(request)
assert num_computed_tokens == 0
blocks = manager.allocate_slots(
request,
num_new_tokens=request.num_tokens,
new_computed_blocks=computed_blocks,
delay_cache_blocks=True,
full_sequence_must_fit=True,
)
assert blocks is not None
assert len(manager.get_block_ids(request.request_id)[0]) == 2
# Empty hashes intentionally mean request-local residency only. Calling
# native cache_blocks() would require a canonical hash for the full prefix
# block and is therefore outside this request's contract.
manager.free(request)
assert manager.block_pool.get_num_free_blocks() == free_before
def test_hashed_request_conforms_to_native_delayed_prefix_commit() -> None:
manager = _manager(enable_caching=True)
prefix_hash = BlockHash(b"prefix".ljust(32, b"\0"))
target_hash = BlockHash(b"target".ljust(32, b"\0"))
request = _request(block_hashes=[prefix_hash, target_hash])
computed_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(request)
assert num_computed_tokens == 0
blocks = manager.allocate_slots(
request,
num_new_tokens=request.num_tokens,
new_computed_blocks=computed_blocks,
delay_cache_blocks=True,
full_sequence_must_fit=True,
)
assert blocks is not None
manager.cache_blocks(request, request.prefix_len)
manager.free(request)
cached_request = _request(
request_id="req/diffusion-kv/cached",
block_hashes=[prefix_hash, target_hash],
)
cached_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(cached_request)
# Only the stable prefix was published. The dynamic target remains a miss.
assert num_computed_tokens == cached_request.prefix_len
assert len(cached_blocks.blocks[0]) == 1
cached_allocation = manager.allocate_slots(
cached_request,
num_new_tokens=cached_request.num_tokens - num_computed_tokens,
num_new_computed_tokens=num_computed_tokens,
new_computed_blocks=cached_blocks,
delay_cache_blocks=True,
full_sequence_must_fit=True,
)
assert cached_allocation is not None
manager.free(cached_request)
@@ -1,290 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.diffusion_kv.metadata import (
DiffusionKVMetadata,
DiffusionKVSequenceMetadata,
)
from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
from vllm_omni.diffusion.sched.interface import (
CachedRequestData,
DiffusionSchedulerOutput,
NewRequestData,
)
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker, WorkerWrapperBase
pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
def make_metadata(request_id: str = "req-0") -> DiffusionKVMetadata:
return DiffusionKVMetadata(
request_id=request_id,
allocation_generation=1,
sequences=(
DiffusionKVSequenceMetadata(
sequence_id=0,
prefix_len=4,
target_len=2,
seq_len=8,
block_ids=([1, 2],),
),
),
)
def make_runner(mode: DiffusionKVCacheMode) -> DiffusionModelRunner:
runner = object.__new__(DiffusionModelRunner)
runner.od_config = SimpleNamespace(diffusion_kv_mode=mode)
return runner
def make_scheduler_output(*new_reqs: NewRequestData) -> DiffusionSchedulerOutput:
return DiffusionSchedulerOutput(
step_id=0,
scheduled_new_reqs=list(new_reqs),
scheduled_cached_reqs=CachedRequestData.make_empty(),
finished_req_ids=set(),
num_running_reqs=len(new_reqs),
num_waiting_reqs=0,
)
def make_executor() -> tuple[MultiprocDiffusionExecutor, list[tuple]]:
executor = object.__new__(MultiprocDiffusionExecutor)
executor.od_config = SimpleNamespace(
enable_distributed_layerwise_offload=False,
dlo_use_allgather=True,
)
executor._ensure_open = lambda: None
calls: list[tuple] = []
def collective_rpc(method, *, args, unique_reply_rank, exec_all_ranks):
calls.append((method, args, unique_reply_rank, exec_all_ranks))
return DiffusionOutput(output=None)
executor.collective_rpc = collective_rpc
return executor, calls
def test_new_request_data_carries_scheduler_allocation_atomically() -> None:
req = SimpleNamespace(request_id="req-0")
state = SimpleNamespace(request_id="req-0", req=req)
metadata = make_metadata()
new_req = NewRequestData.from_state(state, diffusion_kv_metadata=metadata)
assert new_req.req is req
assert new_req.diffusion_kv_metadata is metadata
def test_paged_request_without_metadata_reaches_request_forward() -> None:
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
expected_output = DiffusionOutput(output=None)
runner._execute_request_list = Mock(
return_value=SimpleNamespace(
runner_outputs=[SimpleNamespace(result=expected_output)],
)
)
req = SimpleNamespace(request_id="req-0")
assert runner.execute_model(req) is expected_output
runner._execute_request_list.assert_called_once()
def test_dense_request_rejects_metadata_before_request_forward() -> None:
runner = make_runner(DiffusionKVCacheMode.DENSE_LEGACY)
runner._execute_request_list = Mock()
with pytest.raises(ValueError, match="dense_legacy.*must not carry"):
runner.execute_model(
SimpleNamespace(request_id="req-0"),
diffusion_kv_metadata=make_metadata(),
)
runner._execute_request_list.assert_not_called()
def test_paged_request_rejects_mismatched_metadata_identity() -> None:
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
runner._execute_request_list = Mock()
with pytest.raises(ValueError, match="request mismatch"):
runner.execute_model(
SimpleNamespace(request_id="req-0"),
diffusion_kv_metadata=make_metadata("stale"),
)
runner._execute_request_list.assert_not_called()
def test_request_rpc_rejects_envelope_request_identity_mismatch() -> None:
executor, calls = make_executor()
req = SimpleNamespace(request_id="actual-request-id")
new_req = NewRequestData(
request_id="envelope-id",
req=req,
diffusion_kv_metadata=make_metadata("envelope-id"),
)
with pytest.raises(ValueError, match="request identity mismatch"):
executor.execute_request(make_scheduler_output(new_req))
assert calls == []
def test_batch_path_allows_missing_metadata_before_forward() -> None:
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
expected_output = object()
runner._execute_request_list = Mock(return_value=expected_output)
new_req = NewRequestData(
request_id="req-0",
req=SimpleNamespace(request_id="req-0"),
)
assert runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config) is expected_output
runner._execute_request_list.assert_called_once()
def test_batch_path_rejects_envelope_request_identity_mismatch() -> None:
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
runner._execute_request_list = Mock()
new_req = NewRequestData(
request_id="envelope-id",
req=SimpleNamespace(request_id="actual-request-id"),
diffusion_kv_metadata=make_metadata("envelope-id"),
)
with pytest.raises(ValueError, match="request identity mismatch"):
runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config)
runner._execute_request_list.assert_not_called()
def test_step_path_allows_missing_metadata_before_step_support_check() -> None:
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
runner.pipeline = object()
runner._supports_step_mode = Mock(return_value=False)
new_req = NewRequestData(
request_id="req-0",
req=SimpleNamespace(request_id="req-0"),
)
with pytest.raises(ValueError, match="does not support step execution"):
runner.execute_stepwise(make_scheduler_output(new_req))
runner._supports_step_mode.assert_called_once()
def test_step_path_rejects_envelope_request_identity_mismatch() -> None:
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
runner.pipeline = object()
runner._supports_step_mode = Mock()
new_req = NewRequestData(
request_id="envelope-id",
req=SimpleNamespace(request_id="actual-request-id"),
diffusion_kv_metadata=make_metadata("envelope-id"),
)
with pytest.raises(ValueError, match="request identity mismatch"):
runner.execute_stepwise(make_scheduler_output(new_req))
runner._supports_step_mode.assert_not_called()
def test_dense_request_rpc_keeps_legacy_positional_shape() -> None:
executor, calls = make_executor()
prepared_layout = object()
req = SimpleNamespace(request_id="req-0", prepared_layout=prepared_layout)
new_req = NewRequestData(request_id="req-0", req=req)
result = executor.execute_request(make_scheduler_output(new_req))
assert result.request_ids == ["req-0"]
assert calls == [
(
"execute_model",
(req, executor.od_config, None),
0,
True,
)
]
assert calls[0][1][0].prepared_layout is prepared_layout
def test_request_rpc_sends_attached_metadata_for_only_that_request() -> None:
executor, calls = make_executor()
req_0 = SimpleNamespace(request_id="req-0")
req_1 = SimpleNamespace(request_id="req-1")
metadata_0 = make_metadata("req-0")
metadata_1 = make_metadata("req-1")
new_reqs = (
NewRequestData(request_id="req-0", req=req_0, diffusion_kv_metadata=metadata_0),
NewRequestData(request_id="req-1", req=req_1, diffusion_kv_metadata=metadata_1),
)
result = executor.execute_request(make_scheduler_output(*new_reqs))
assert result.request_ids == ["req-0", "req-1"]
assert [call[1] for call in calls] == [
(req_0, executor.od_config, None, metadata_0),
(req_1, executor.od_config, None, metadata_1),
]
def test_worker_wrapper_only_extends_paged_request_rpc() -> None:
calls: list[tuple] = []
class Worker:
def execute_model(self, req, od_config, **kwargs):
calls.append((req, od_config, kwargs))
return DiffusionOutput(output=None)
wrapper = object.__new__(WorkerWrapperBase)
wrapper.worker = Worker()
req = SimpleNamespace(request_id="req-0")
od_config = SimpleNamespace()
metadata = make_metadata()
wrapper.execute_model(req, od_config)
wrapper.execute_model(req, od_config, diffusion_kv_metadata=metadata)
assert calls == [
(req, od_config, {"kv_prefetch_job": None}),
(
req,
od_config,
{
"kv_prefetch_job": None,
"diffusion_kv_metadata": metadata,
},
),
]
def test_dense_worker_call_does_not_extend_model_runner_signature() -> None:
calls: list[tuple] = []
class DenseModelRunner:
def execute_model(self, req, kv_prefetch_job=None):
calls.append((req, kv_prefetch_job))
return DiffusionOutput(output=None)
worker = object.__new__(DiffusionWorker)
worker.model_runner = DenseModelRunner()
worker.lora_manager = None
worker._get_profiler = lambda: None
req = SimpleNamespace(request_id="req-0")
worker.execute_model(req, SimpleNamespace())
assert calls == [(req, None)]
@@ -1,447 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import pytest
import torch
from transformers.utils.generic import ModelOutput
import vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 as pipeline_module
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_tokenizer import TokenizerEncodeOutput
from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
ImageInfo,
JointImageInfo,
)
from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
HunyuanImage3Pipeline,
get_hunyuan_image_3_pre_process_func,
)
from vllm_omni.diffusion.models.hunyuan_image3.request_layout import (
HunyuanPreparedLayout,
build_hunyuan_diffusion_kv_requests,
extract_hunyuan_prompt_inputs,
hunyuan_num_image_tokens,
normalize_hunyuan_cot_text,
prepare_hunyuan_layout,
)
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def test_dense_legacy_preprocess_does_not_initialize_layout_tokenizer(monkeypatch) -> None:
hf_config = SimpleNamespace(vae_downsample_factor=(8, 8), patch_size=2)
image_processor = SimpleNamespace(vision_encoder_processor=SimpleNamespace(patch_size=1))
monkeypatch.setattr(pipeline_module, "get_config", lambda *_args, **_kwargs: hf_config)
monkeypatch.setattr(pipeline_module, "HunyuanImage3ImageProcessor", lambda _config: image_processor)
monkeypatch.setattr(
pipeline_module,
"TokenizerWrapper",
lambda _model: pytest.fail("dense_legacy must not initialize the Diffusion KV tokenizer"),
)
get_hunyuan_image_3_pre_process_func(
SimpleNamespace(model="model", diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
)
class _FakeTokenizerWrapper:
def __init__(self, *, prefix_lens: list[int]) -> None:
self.prefix_lens = prefix_lens
self.calls: list[dict] = []
self.pad_token_id = 0
self.eos_token_id = 2
self.boi_token_id = 3
self.end_recaption_token_id = 4
self.end_answer_token_id = 5
self.special_token_map = {f"<img_ratio_{index}>": 100 + index for index in range(33)}
def apply_chat_template(self, **kwargs):
self.calls.append(kwargs)
rows = len(self.prefix_lens)
has_cond_image = kwargs.get("batch_cond_image_info") is not None
seq_lens = [prefix_len + 20 for prefix_len in self.prefix_lens]
padded_seq_len = max(seq_lens)
positions = torch.arange(padded_seq_len)
joint_image_slices = [[slice(1, 3), slice(3, 5)] if has_cond_image else [] for _ in range(rows)]
gen_image_slices = [[slice(length + 1, length + 17)] for length in self.prefix_lens]
all_image_slices = [joint + generated for joint, generated in zip(joint_image_slices, gen_image_slices)]
cond_vae_image_mask = (
torch.stack([(positions >= 1) & (positions < 3) for _ in range(rows)]) if has_cond_image else None
)
cond_vit_image_mask = (
torch.stack([(positions >= 3) & (positions < 5) for _ in range(rows)]) if has_cond_image else None
)
sections = []
for _ in range(rows):
row_sections = []
if has_cond_image:
row_sections.append(
{
"type": "joint_image",
"token_height": [1, 1],
"token_width": [2, 2],
}
)
row_sections.append({"type": "gen_image", "token_height": 4, "token_width": 4})
sections.append(row_sections)
return {
"output": TokenizerEncodeOutput(
tokens=torch.arange(rows * padded_seq_len, dtype=torch.long).reshape(rows, padded_seq_len),
gen_timestep_scatter_index=torch.tensor(self.prefix_lens, dtype=torch.long).reshape(rows, 1),
cond_timestep_scatter_index=(torch.zeros(rows, 1, dtype=torch.long) if has_cond_image else None),
all_image_slices=all_image_slices,
gen_image_mask=torch.stack(
[(positions >= length + 1) & (positions < length + 17) for length in self.prefix_lens]
),
cond_vae_image_mask=cond_vae_image_mask,
cond_vit_image_mask=cond_vit_image_mask,
joint_image_slices=joint_image_slices,
gen_image_slices=gen_image_slices,
real_pos=torch.tensor(seq_lens, dtype=torch.long).reshape(rows, 1),
),
"sections": sections,
}
class _FakeImageProcessor:
def __init__(self, image_token_length: int = 16) -> None:
self.image_token_length = image_token_length
self.image_sizes: list[tuple[int, int]] = []
self.vision_encoder_processor = SimpleNamespace(patch_size=1)
def build_image_info(self, image_size):
self.image_sizes.append(image_size)
return ImageInfo(
image_type="gen_image",
image_width=image_size[1],
image_height=image_size[0],
token_width=4,
token_height=4,
image_token_length=self.image_token_length,
base_size=1024,
ratio_index=0,
)
def _components(prefix_lens: list[int]):
return _FakeTokenizerWrapper(prefix_lens=prefix_lens), _FakeImageProcessor()
def _prepare(
request: OmniDiffusionRequest,
tokenizer: _FakeTokenizerWrapper,
image_processor: _FakeImageProcessor,
) -> HunyuanPreparedLayout:
prepared_layout = prepare_hunyuan_layout(
request,
tokenizer_wrapper=tokenizer,
image_processor=image_processor,
generation_config=SimpleNamespace(sequence_template="instruct", drop_think=False),
image_base_size=1024,
)
request.prepared_layout = prepared_layout
return prepared_layout
def _request(
*,
guidance_scale: float,
prompt="draw a cat",
request_id: str = "req",
) -> OmniDiffusionRequest:
return OmniDiffusionRequest(
prompt=prompt,
sampling_params=OmniDiffusionSamplingParams(
height=768,
width=1024,
guidance_scale=guidance_scale,
num_inference_steps=4,
),
request_id=request_id,
)
def test_builds_kv_request_lengths_without_model_execution() -> None:
tokenizer, image_processor = _components([12])
request = _request(guidance_scale=1.0)
prepared_layout = _prepare(request, tokenizer, image_processor)
kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
assert isinstance(request.prepared_layout, HunyuanPreparedLayout)
assert all(isinstance(item, DiffusionKVRequest) for item in kv_requests)
assert [(item.prefix_len, item.target_len, item.seq_len) for item in kv_requests] == [(12, 17, 32)]
assert image_processor.image_sizes == [(768, 1024)]
assert tokenizer.calls[0]["sequence_template"] == "instruct"
assert tokenizer.calls[0]["cfg_factor"] == 1
assert kv_requests[0].block_hashes == []
assert kv_requests[0].kv_contexts == ()
assert kv_requests[0].skip_reading_prefix_cache is True
def test_builds_one_kv_request_per_cfg_row() -> None:
tokenizer, image_processor = _components([12, 14])
request = _request(guidance_scale=5.0)
prepared_layout = _prepare(request, tokenizer, image_processor)
kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
assert [item.sequence_id for item in kv_requests] == [0, 1]
assert [item.prefix_len for item in kv_requests] == [12, 14]
assert [item.seq_len for item in kv_requests] == [32, 34]
assert tokenizer.calls[0]["cfg_factor"] == 2
def _reference_image() -> JointImageInfo:
return JointImageInfo(
vae_image_info=ImageInfo(
image_type="vae",
token_width=8,
token_height=8,
image_token_length=64,
),
vision_image_info=ImageInfo(
image_type="siglip2",
token_width=4,
token_height=4,
image_token_length=16,
),
vision_encoder_kwargs={
"spatial_shapes": torch.tensor([4, 4]),
"pixel_attention_mask": torch.ones(4, dtype=torch.bool),
},
)
def test_passes_preprocessed_reference_image_geometry_to_tokenizer() -> None:
tokenizer, image_processor = _components([20])
joint_image = _reference_image()
prompt = {
"prompt": "edit this image",
"additional_information": {"batch_cond_image_info": [joint_image]},
}
request = _request(guidance_scale=1.0, prompt=prompt)
prepared_layout = _prepare(request, tokenizer, image_processor)
kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
assert tokenizer.calls[0]["batch_cond_image_info"] == [[joint_image]]
assert kv_requests[0].prefix_len == 20
assert kv_requests[0].target_len == 17
assert kv_requests[0].seq_len == 40
assert kv_requests[0].kv_contexts == ()
def _assert_nested_equal(actual, expected) -> None:
if isinstance(expected, torch.Tensor):
torch.testing.assert_close(actual, expected)
elif isinstance(expected, dict):
assert actual.keys() == expected.keys()
for key in expected:
_assert_nested_equal(actual[key], expected[key])
elif isinstance(expected, list | tuple):
assert isinstance(actual, type(expected))
assert len(actual) == len(expected)
for actual_item, expected_item in zip(actual, expected):
_assert_nested_equal(actual_item, expected_item)
else:
assert actual == expected
@pytest.mark.parametrize(
("prefix_lens", "guidance_scale", "with_reference_image"),
[
pytest.param([12], 1.0, False, id="cfg-1"),
pytest.param([12, 14], 5.0, False, id="cfg-2"),
pytest.param([20], 1.0, True, id="cfg-1-reference-image"),
pytest.param([20, 22], 5.0, True, id="cfg-2-reference-image"),
],
)
def test_prepared_model_inputs_match_local_tokenization(
monkeypatch,
prefix_lens: list[int],
guidance_scale: float,
with_reference_image: bool,
) -> None:
tokenizer, image_processor = _components(prefix_lens)
prompt: str | dict = "draw a cat"
if with_reference_image:
prompt = {
"prompt": "edit this image",
"additional_information": {"batch_cond_image_info": [_reference_image()]},
}
request = _request(guidance_scale=guidance_scale, prompt=prompt)
prepared_layout = _prepare(request, tokenizer, image_processor)
prompts, cot_texts, system_prompt, batch_cond_image_info, bot_task = extract_hunyuan_prompt_inputs(
[request.prompt],
request.sampling_params.extra_args or {},
request_id=request.request_id,
allow_cond_image=True,
)
cot_text = (
[normalize_hunyuan_cot_text(text) for text in cot_texts]
if any(text is not None for text in cot_texts)
else None
)
pipeline = object.__new__(HunyuanImage3Pipeline)
pipeline._tkwrapper = tokenizer
pipeline.image_processor = image_processor
pipeline.generation_config = SimpleNamespace(sequence_template="instruct", drop_think=False)
pipeline.config = SimpleNamespace(
image_base_size=1024,
attention_head_dim=2,
rope_theta=10000.0,
)
monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda _self: torch.device("cpu")))
def fake_encode_cond_image(batch_cond_image_info_list, cfg_factor=1, generator=None):
del generator
rows = len(batch_cond_image_info_list) * cfg_factor
return (
torch.arange(rows * 4, dtype=torch.float32).reshape(rows, 1, 2, 2),
torch.arange(rows, dtype=torch.float32),
torch.arange(rows * 3, dtype=torch.float32).reshape(rows, 3),
)
pipeline._encode_cond_image = fake_encode_cond_image
def fake_build_batch_2d_rope(*, image_infos, seq_len, **_kwargs):
geometry_sum = sum(
int(image_slice.start or 0) + int(image_slice.stop or 0) + height + width
for row in image_infos
for image_slice, (height, width) in row
)
values = torch.tensor([len(image_infos), seq_len, geometry_sum], dtype=torch.float32)
return values, values + 1
monkeypatch.setattr(pipeline_module, "build_batch_2d_rope", fake_build_batch_2d_rope)
common_kwargs = dict(
prompt=prompts,
cot_text=cot_text,
system_prompt=system_prompt,
mode="gen_image",
guidance_scale=guidance_scale,
image_size=(request.sampling_params.height, request.sampling_params.width),
generator=[torch.Generator().manual_seed(0)],
batch_cond_image_info=batch_cond_image_info,
bot_task=bot_task,
)
local_inputs = pipeline.prepare_model_inputs(**common_kwargs)
prepared_inputs = pipeline.prepare_model_inputs(**common_kwargs, prepared_layout=prepared_layout)
comparable_fields = (
"input_ids",
"position_ids",
"custom_pos_emb",
"image_mask",
"gen_timestep_scatter_index",
"cond_vae_images",
"cond_timestep",
"cond_vae_image_mask",
"cond_vit_images",
"cond_vit_image_mask",
"vit_kwargs",
"cond_timestep_scatter_index",
)
for field in comparable_fields:
_assert_nested_equal(prepared_inputs[field], local_inputs[field])
tokenizer_fields = (
"tokens",
"gen_image_mask",
"cond_vae_image_mask",
"cond_vit_image_mask",
"gen_timestep_scatter_index",
"cond_timestep_scatter_index",
"real_pos",
"all_image_slices",
"joint_image_slices",
"gen_image_slices",
)
for field in tokenizer_fields:
_assert_nested_equal(
getattr(prepared_inputs["tokenizer_output"], field),
getattr(local_inputs["tokenizer_output"], field),
)
local_mask = pipeline._prepare_attention_mask_for_generation(
local_inputs["input_ids"],
pipeline.generation_config,
local_inputs,
)
prepared_mask = pipeline._prepare_attention_mask_for_generation(
prepared_inputs["input_ids"],
pipeline.generation_config,
prepared_inputs,
)
torch.testing.assert_close(prepared_mask, local_mask)
assert prepared_inputs["full_attn_spans"] == local_inputs["full_attn_spans"]
num_image_tokens = hunyuan_num_image_tokens(prepared_layout.generated_image_info)
local_inputs.update(attention_mask=local_mask, num_image_tokens=num_image_tokens)
prepared_inputs.update(attention_mask=prepared_mask, num_image_tokens=num_image_tokens)
local_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), local_inputs)
prepared_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), prepared_inputs)
for field in ("position_ids", "attention_mask", "gen_timestep_scatter_index", "full_attn_spans"):
_assert_nested_equal(prepared_step_inputs[field], local_step_inputs[field])
def test_rejects_tokenizer_cfg_row_mismatch() -> None:
tokenizer, image_processor = _components([12])
request = _request(guidance_scale=5.0)
prepared_layout = _prepare(request, tokenizer, image_processor)
with pytest.raises(ValueError, match="sequence count does not match"):
build_hunyuan_diffusion_kv_requests(request, prepared_layout)
def test_sender_endpoint_does_not_change_local_kv_request() -> None:
tokenizer, image_processor = _components([12])
request = _request(guidance_scale=1.0)
request.kv_sender_info = {"host": "127.0.0.1", "zmq_port": 5000}
prepared_layout = _prepare(request, tokenizer, image_processor)
kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
assert kv_requests[0].seq_len == 32
assert kv_requests[0].num_computed_tokens == 0
def test_paged_preprocess_attaches_layout_and_scheduler_kv_requests(monkeypatch) -> None:
tokenizer, image_processor = _components([12])
hf_config = SimpleNamespace(
vae_downsample_factor=(8, 8),
patch_size=2,
image_base_size=1024,
)
monkeypatch.setattr(pipeline_module, "get_config", lambda *_args, **_kwargs: hf_config)
monkeypatch.setattr(pipeline_module, "HunyuanImage3ImageProcessor", lambda _config: image_processor)
monkeypatch.setattr(pipeline_module, "TokenizerWrapper", lambda _model: tokenizer)
monkeypatch.setattr(
pipeline_module.GenerationConfig,
"from_pretrained",
lambda _model: SimpleNamespace(sequence_template="instruct", drop_think=False),
)
preprocess = get_hunyuan_image_3_pre_process_func(
SimpleNamespace(model="model", diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
)
request = _request(guidance_scale=1.0)
prepared_request = preprocess(request)
assert prepared_request is request
assert isinstance(request.prepared_layout, HunyuanPreparedLayout)
assert request.diffusion_kv_requests is not None
assert [(item.prefix_len, item.target_len, item.seq_len) for item in request.diffusion_kv_requests] == [
(12, 17, 32)
]
assert len(tokenizer.calls) == 1
@@ -7,10 +7,7 @@ import pytest
import torch
import vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 as hy3_module
import vllm_omni.diffusion.models.hunyuan_image3.request_layout as hy3_layout_module
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_tokenizer import TokenizerEncodeOutput
from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import ImageInfo
from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
_STEP_AR_KV,
_STEP_CFG_FACTOR,
@@ -21,7 +18,6 @@ from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
_STEP_PROMPT_KV,
HunyuanImage3Pipeline,
)
from vllm_omni.diffusion.models.hunyuan_image3.request_layout import HunyuanPreparedLayout
from vllm_omni.diffusion.worker.input_batch import InputBatch
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.diffusion.worker.utils import StepRequestState
@@ -80,72 +76,6 @@ def _sampling_params(**extra_args):
)
def _prepared_layout() -> HunyuanPreparedLayout:
return HunyuanPreparedLayout(
tokenizer_output=TokenizerEncodeOutput(
tokens=torch.arange(21).reshape(1, 21),
gen_image_mask=torch.ones(1, 21, dtype=torch.bool),
gen_timestep_scatter_index=torch.tensor([[4]]),
all_image_slices=[[slice(5, 21)]],
joint_image_slices=[[]],
gen_image_slices=[[slice(5, 21)]],
real_pos=torch.tensor([[21]]),
),
rope_image_info=[[(slice(5, 21), (4, 4))]],
generated_image_info=ImageInfo(
image_type="gen_image",
image_width=512,
image_height=512,
token_width=4,
token_height=4,
image_token_length=16,
),
)
def test_prepare_model_inputs_reuses_prepared_layout(monkeypatch):
pipeline = _pipeline()
monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda self: torch.device("cpu")))
rope_kwargs = {}
def fake_build_batch_2d_rope(**kwargs):
rope_kwargs.update(kwargs)
return torch.zeros(1), torch.zeros(1)
monkeypatch.setattr(hy3_module, "build_batch_2d_rope", fake_build_batch_2d_rope)
def fail_apply_chat_template(**_kwargs):
pytest.fail("prepared Hunyuan execution must not apply the chat template again")
pipeline._tkwrapper = SimpleNamespace(
apply_chat_template=fail_apply_chat_template,
eos_token_id=2,
boi_token_id=3,
end_recaption_token_id=4,
end_answer_token_id=5,
)
pipeline.config = SimpleNamespace(
attention_head_dim=2,
rope_theta=10000.0,
)
prepared = _prepared_layout()
model_inputs = pipeline.prepare_model_inputs(
prompt="prompt",
mode="gen_image",
guidance_scale=1.0,
image_size=(512, 512),
device=torch.device("cpu"),
generator=[torch.Generator().manual_seed(0)],
prepared_layout=prepared,
)
assert model_inputs["tokenizer_output"] is prepared.tokenizer_output
assert model_inputs["batch_gen_image_info"] == [prepared.generated_image_info]
assert rope_kwargs["image_infos"] is prepared.rope_image_info
torch.testing.assert_close(model_inputs["input_ids"], prepared.tokenizer_output.tokens)
def test_hunyuan_step_group_key_ignores_step_index_for_later_steps():
pipeline = _pipeline()
states = [_state("req-0", 1), _state("req-1", 3)]
@@ -201,13 +131,12 @@ def test_prepare_encode_preserves_normal_hunyuan_bot_task_semantics(
captured.update(kwargs)
raise RuntimeError("stop after prepare_model_inputs")
monkeypatch.setattr(hy3_layout_module, "get_system_prompt", fake_get_system_prompt)
monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt)
pipeline.prepare_model_inputs = fake_prepare_model_inputs
state = StepRequestState(
request_id="req-bot-task",
sampling=sampling,
prompt=prompt_item,
prepared_layout=_prepared_layout(),
)
with pytest.raises(RuntimeError, match="stop after prepare_model_inputs"):
@@ -215,7 +144,6 @@ def test_prepare_encode_preserves_normal_hunyuan_bot_task_semantics(
assert captured["bot_task"] == expected_model_bot_task
assert captured["system_prompt_bot_task"] == expected_system_bot_task
assert captured["prepared_layout"] is state.prepared_layout
def test_forward_uses_same_hunyuan_bot_task_semantics(monkeypatch):
@@ -231,22 +159,23 @@ def test_forward_uses_same_hunyuan_bot_task_semantics(monkeypatch):
captured.update(kwargs)
raise RuntimeError("stop after prepare_model_inputs")
monkeypatch.setattr(hy3_layout_module, "get_system_prompt", fake_get_system_prompt)
monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt)
pipeline.prepare_model_inputs = fake_prepare_model_inputs
request = SimpleNamespace(
request_id="req-forward-bot-task",
sampling_params=_sampling_params(bot_task="think_recaption", use_system_prompt="dynamic"),
prompt={"prompt": "prompt", "bot_task": "vanilla"},
prepared_layout=_prepared_layout(),
req = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="req-forward-bot-task",
sampling_params=_sampling_params(bot_task="think_recaption", use_system_prompt="dynamic"),
prompt={"prompt": "prompt", "bot_task": "vanilla"},
)
]
)
req = DiffusionRequestBatch(requests=[request])
with pytest.raises(RuntimeError, match="stop after prepare_model_inputs"):
pipeline.forward(req)
assert captured["bot_task"] == "think"
assert captured["system_prompt_bot_task"] == "think"
assert captured["prepared_layout"] is request.prepared_layout
def test_grouped_denoise_rejects_non_sdpa_attention_backend():
@@ -1,133 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import numpy as np
import pytest
import torch
from PIL import Image
from torch import nn
from vllm_omni.diffusion.models.lingbot_video.image_condition import (
apply_clean_prefix,
geometry_align_image,
prepare_ti2v_image_condition,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def _gradient_image(width: int, height: int) -> Image.Image:
x = np.linspace(0, 255, width, dtype=np.uint8)[None, :]
y = np.linspace(0, 255, height, dtype=np.uint8)[:, None]
rgb = np.stack(
[
np.broadcast_to(x, (height, width)),
np.broadcast_to(y, (height, width)),
np.full((height, width), 127, dtype=np.uint8),
],
axis=-1,
)
return Image.fromarray(rgb, mode="RGB")
@pytest.mark.parametrize("source_size", [(160, 64), (64, 160), (96, 96)])
def test_geometry_align_image_supports_source_aspect_ratios(source_size):
pixel = geometry_align_image(_gradient_image(*source_size), 64, 96)
assert pixel.shape == (1, 3, 1, 64, 96)
assert pixel.dtype == torch.float32
assert 0 <= pixel.min() <= pixel.max() <= 1
@pytest.mark.parametrize(
("source_size", "target_size"),
[((160, 64), (64, 96)), ((64, 160), (96, 64))],
)
def test_geometry_align_image_center_crops_without_stretching(source_size, target_size):
width, height = source_size
target_height, target_width = target_size
rgb = np.full((height, width, 3), (255, 0, 0), dtype=np.uint8)
if width > target_width:
left = (width - target_width) // 2
rgb[:, left : left + target_width] = (0, 255, 0)
else:
top = (height - target_height) // 2
rgb[top : top + target_height] = (0, 255, 0)
pixel = geometry_align_image(
Image.fromarray(rgb, mode="RGB"),
target_height,
target_width,
)
assert pixel.shape == (1, 3, 1, target_height, target_width)
expected_green = torch.tensor((0.0, 1.0, 0.0)).view(1, 3, 1, 1, 1)
assert torch.equal(pixel, expected_green.expand_as(pixel))
class _LatentDistribution:
def __init__(self, latent: torch.Tensor):
self.latent = latent
def sample(self, generator=None):
del generator
return self.latent
class _CapturingVAE(nn.Module):
def __init__(self):
super().__init__()
self.anchor = nn.Parameter(torch.zeros(()))
self.config = SimpleNamespace(
latents_mean=[0.0, 0.0, 0.0],
latents_std=[1.0, 1.0, 1.0],
)
self.encoded_pixel = None
def encode(self, pixel):
self.encoded_pixel = pixel.detach().clone()
return SimpleNamespace(latent_dist=_LatentDistribution(pixel))
def test_ti2v_vlm_and_vae_paths_share_geometry_aligned_image():
vae = _CapturingVAE()
condition = prepare_ti2v_image_condition(
_gradient_image(160, 64),
height=64,
width=96,
vae=vae,
vision_patch_size=16,
device=torch.device("cpu"),
)
vae_pixel = (vae.encoded_pixel + 1.0) / 2.0
vae_image = vae_pixel[0, :, 0].permute(1, 2, 0).mul(255).byte().numpy()
assert condition.vlm_image.size == (96, 64)
np.testing.assert_allclose(
np.asarray(condition.vlm_image),
vae_image,
rtol=0,
atol=1,
)
def test_apply_clean_prefix_only_overwrites_condition_frames():
latents = torch.arange(20, dtype=torch.float32).reshape(1, 1, 5, 2, 2)
original_tail = latents[:, :, 2:].clone()
clean = torch.full((1, 1, 2, 2, 2), -3.0)
returned = apply_clean_prefix(latents, clean)
assert returned is latents
assert torch.equal(latents[:, :, :2], clean)
assert torch.equal(latents[:, :, 2:], original_tail)
def test_apply_clean_prefix_rejects_incompatible_shapes():
with pytest.raises(ValueError, match="shape must match"):
apply_clean_prefix(
torch.zeros(1, 2, 3, 4, 4),
torch.zeros(1, 3, 1, 4, 4),
)
@@ -6,10 +6,9 @@ from types import SimpleNamespace
import numpy as np
import pytest
import torch
from PIL import Image
from torch import nn
from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -23,8 +22,6 @@ def _make_pipeline():
nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
pipeline.default_negative_prompt = "default negative"
pipeline.default_image_negative_prompt = "default image negative"
pipeline.img_prompt_template = "<image>"
pipeline.od_config = SimpleNamespace(flow_shift=None)
return pipeline
@@ -40,24 +37,17 @@ def test_lingbot_video_pipeline_import_and_registry():
LingBotVideoPipeline,
LingBotVideoTransformer3DModel,
get_lingbot_video_post_process_func,
get_lingbot_video_pre_process_func,
)
from vllm_omni.diffusion.registry import (
_DIFFUSION_MODELS,
_DIFFUSION_POST_PROCESS_FUNCS,
_DIFFUSION_PRE_PROCESS_FUNCS,
)
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS, _DIFFUSION_POST_PROCESS_FUNCS
assert LingBotVideoPipeline is not None
assert LingBotVideoTransformer3DModel is not None
assert get_lingbot_video_pre_process_func is not None
assert get_lingbot_video_post_process_func is not None
assert _DIFFUSION_MODELS["LingBotVideoPipeline"] == (
"lingbot_video",
"pipeline_lingbot_video",
"LingBotVideoPipeline",
)
assert _DIFFUSION_PRE_PROCESS_FUNCS["LingBotVideoPipeline"] == "get_lingbot_video_pre_process_func"
assert _DIFFUSION_POST_PROCESS_FUNCS["LingBotVideoPipeline"] == "get_lingbot_video_post_process_func"
@@ -68,7 +58,6 @@ def test_component_discovery_declarations():
assert LingBotVideoPipeline._encoder_modules == ["text_encoder"]
assert LingBotVideoPipeline._vae_modules == ["vae"]
assert LingBotVideoPipeline.supports_step_execution is False
assert LingBotVideoPipeline.max_outputs_per_prompt == 1
def test_extra_body_params_include_video_flow_shift_alias():
@@ -80,47 +69,17 @@ def test_extra_body_params_include_video_flow_shift_alias():
assert "batch_cfg" in params
def test_preprocess_preserves_geometry_and_converts_rgb(tmp_path):
from vllm_omni.diffusion.models.lingbot_video import (
get_lingbot_video_pre_process_func,
)
def test_extract_prompt_string_and_mapping():
from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import _extract_prompt
image_path = tmp_path / "reference.png"
Image.new("L", (48, 32), color=127).save(image_path)
request = OmniDiffusionRequest(
prompt={
"prompt": "a robot walks",
"modalities": ["video"],
"multi_modal_data": {"image": str(image_path)},
},
sampling_params=OmniDiffusionSamplingParams(),
request_id="lingbot-preprocess-test",
)
string_req = SimpleNamespace(prompt="a robot waves")
assert _extract_prompt(string_req) == ("a robot waves", None)
result = get_lingbot_video_pre_process_func(SimpleNamespace())(request)
mapping_req = SimpleNamespace(prompt={"prompt": "a robot walks", "negative_prompt": "low quality"})
assert _extract_prompt(mapping_req) == ("a robot walks", "low quality")
image = result.prompt["multi_modal_data"]["image"]
assert image.mode == "RGB"
assert image.size == (48, 32)
def test_preprocess_marks_internal_image_warmup_as_video():
from vllm_omni.diffusion.models.lingbot_video import (
get_lingbot_video_pre_process_func,
)
request = OmniDiffusionRequest(
prompt={
"prompt": "dummy run",
"multi_modal_data": {"image": Image.new("RGB", (32, 32))},
},
sampling_params=OmniDiffusionSamplingParams(),
request_id=DUMMY_DIFFUSION_REQUEST_ID,
)
result = get_lingbot_video_pre_process_func(SimpleNamespace())(request)
assert result.prompt["modalities"] == ["video"]
with pytest.raises(TypeError, match="Unsupported LingBot prompt type"):
_extract_prompt(SimpleNamespace(prompt=["not", "supported"]))
def test_check_inputs_accepts_t2v_shapes_and_rejects_invalid_values():
@@ -145,59 +104,6 @@ def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
array = post(frames, SimpleNamespace(output_type="np"))
assert isinstance(array, np.ndarray)
assert array.shape == (2, 4, 4, 3)
video = post({"video": frames}, SimpleNamespace(output_type="pt"))
assert set(video) == {"video"}
assert video["video"] is frames
image = post(
{"image": torch.ones(4, 4, 3)},
SimpleNamespace(output_type="pt"),
)
assert set(image) == {"image"}
assert isinstance(image["image"], np.ndarray)
assert image["image"].shape == (4, 4, 3)
latent = torch.ones(1, 4, 1, 2, 2)
latent_output = post({"image": latent}, SimpleNamespace(output_type="latent"))
assert set(latent_output) == {"image"}
assert latent_output["image"] is latent
def test_postprocess_preserves_image_contract_for_serving():
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
from vllm_omni.diffusion.output_formatter import (
format_diffusion_outputs,
normalize_diffusion_postprocess_output,
)
from vllm_omni.entrypoints.openai.api_server import _extract_images_from_result
sampling = OmniDiffusionSamplingParams(
num_inference_steps=1,
output_type="pt",
resolution=64,
)
request = OmniDiffusionRequest(
prompt={"prompt": "a robot", "modalities": ["image"]},
sampling_params=sampling,
request_id="lingbot-image-contract-test",
)
raw_output = {"image": torch.ones(8, 12, 3)}
processed = get_lingbot_video_post_process_func(SimpleNamespace())(raw_output, sampling)
normalized = normalize_diffusion_postprocess_output(processed)
assert normalized.primary_key == "image"
[result] = format_diffusion_outputs(
request=request,
od_config=SimpleNamespace(model_class_name="LingBotVideoPipeline"),
diffusion_output=DiffusionOutput(output=raw_output),
output_data=raw_output,
postprocess_output=normalized,
)
assert len(result.images) == 1
assert isinstance(result.images[0], np.ndarray)
endpoint_images = _extract_images_from_result(result)
assert len(endpoint_images) == 1
assert endpoint_images[0].size == (12, 8)
def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
@@ -210,11 +116,7 @@ def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
pipeline._generate = fake_generate
req = _make_request_batch(
{
"prompt": "a robot arm",
"negative_prompt": "bad hands",
"modalities": ["video"],
},
{"prompt": "a robot arm", "negative_prompt": "bad hands"},
height=192,
width=320,
num_frames=9,
@@ -227,12 +129,10 @@ def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
out = pipeline.forward(req)
assert torch.equal(out.output["video"], torch.zeros(9, 192, 320, 3))
assert torch.equal(out.output, torch.zeros(9, 192, 320, 3))
assert len(calls) == 1
call = calls[0]
assert call["prompt"] == "a robot arm"
assert call["mode"].value == "t2v"
assert call["input_image"] is None
assert call["negative_prompt"] == "bad hands"
assert call["height"] == 192
assert call["width"] == 320
@@ -294,222 +194,6 @@ def test_forward_prefers_shift_over_flow_shift_and_defaults_negative_prompt():
assert calls[0]["negative_prompt"] == "default negative"
@pytest.mark.parametrize(
("prompt", "num_frames", "expected_mode", "output_key", "default_negative"),
[
(
{"prompt": "a still robot", "modalities": ["image"]},
1,
"t2i",
"image",
"default image negative",
),
(
{
"prompt": "a moving robot",
"modalities": ["video"],
"multi_modal_data": {"image": Image.new("RGB", (48, 32), color="blue")},
},
9,
"ti2v",
"video",
"default negative",
),
],
)
def test_forward_dispatches_generation_mode_and_output(
prompt,
num_frames,
expected_mode,
output_key,
default_negative,
):
pipeline = _make_pipeline()
calls = []
def fake_generate(**kwargs):
calls.append(kwargs)
return torch.zeros(192, 320, 3) if output_key == "image" else torch.zeros(num_frames, 192, 320, 3)
pipeline._generate = fake_generate
output = pipeline.forward(
_make_request_batch(
prompt,
height=192,
width=320,
num_frames=num_frames,
num_inference_steps=2,
guidance_scale=3.0,
seed=42,
output_type="pt",
)
)
assert set(output.output) == {output_key}
assert calls[0]["mode"].value == expected_mode
assert calls[0]["negative_prompt"] == default_negative
if expected_mode == "ti2v":
assert isinstance(calls[0]["input_image"], Image.Image)
def test_batched_cfg_pads_image_conditioned_embeddings_and_masks():
from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import (
_batch_cfg_prompt_inputs,
)
positive = torch.ones(1, 7, 4)
positive_mask = torch.ones(1, 7, dtype=torch.long)
negative = torch.full((1, 9, 4), 2.0)
negative_mask = torch.ones(1, 9, dtype=torch.long)
embeds, mask = _batch_cfg_prompt_inputs(
positive,
positive_mask,
negative,
negative_mask,
)
assert embeds.shape == (2, 9, 4)
assert mask.shape == (2, 9)
assert torch.equal(mask[0], torch.tensor([1, 1, 1, 1, 1, 1, 1, 0, 0]))
assert torch.equal(mask[1], torch.ones(9, dtype=torch.long))
class _RecordingTransformer(nn.Module):
def __init__(self):
super().__init__()
self.anchor = nn.Parameter(torch.zeros(()))
self.config = SimpleNamespace(in_channels=1)
self.latent_inputs = []
def forward(
self,
latents,
timestep,
prompt_embeds,
*,
encoder_attention_mask,
return_dict,
):
del timestep, prompt_embeds, encoder_attention_mask, return_dict
self.latent_inputs.append(latents.detach().clone())
return (torch.zeros_like(latents),)
class _CorruptingScheduler:
sigma_max = 1.0
sigma_min = 0.0
def set_timesteps(self, num_inference_steps, *, device, shift, **kwargs):
del num_inference_steps, shift, kwargs
self.timesteps = torch.tensor([1000.0, 500.0], device=device)
def step(
self,
noise_pred,
timestep,
latents,
*,
return_dict,
generator,
):
del noise_pred, timestep, return_dict, generator
return (latents + 7.0,)
@pytest.mark.parametrize(
("guidance_scale", "batch_cfg", "expected_prompts", "expected_transformer_calls"),
[
pytest.param(2.0, False, ["positive", "negative"], 4, id="sequential-cfg"),
pytest.param(2.0, True, ["positive", "negative"], 2, id="batched-cfg"),
pytest.param(1.0, False, ["positive"], 2, id="cfg-disabled"),
],
)
def test_ti2v_reinjects_clean_prefix_across_cfg_modes(
guidance_scale,
batch_cfg,
expected_prompts,
expected_transformer_calls,
):
from vllm_omni.diffusion.models.lingbot_video import (
LingBotGenerationMode,
LingBotImageCondition,
)
pipeline = _make_pipeline()
pipeline.transformer = _RecordingTransformer()
pipeline.scheduler = _CorruptingScheduler()
pipeline.progress_bar = lambda timesteps: timesteps
pipeline.prepare_latents = lambda *args, **kwargs: torch.zeros(1, 1, 3, 1, 1)
vlm_image = Image.new("RGB", (32, 32), color="red")
condition = LingBotImageCondition(
vlm_image=vlm_image,
clean_latent=torch.full((1, 1, 1, 1, 1), 5.0),
)
pipeline.prepare_ti2v_image_condition = lambda *args, **kwargs: condition
encode_calls = []
def fake_encode(prompt, *, images, device):
del device
encode_calls.append((prompt, images))
return torch.ones(1, 2, 4), torch.ones(1, 2, dtype=torch.long)
pipeline.encode_prompt = fake_encode
latents = pipeline._generate(
prompt="positive",
mode=LingBotGenerationMode.TI2V,
input_image=vlm_image,
negative_prompt="negative",
height=16,
width=16,
num_frames=9,
num_inference_steps=2,
guidance_scale=guidance_scale,
shift=3.0,
output_type="latent",
batch_cfg=batch_cfg,
)
assert [call[0] for call in encode_calls] == expected_prompts
assert all(call[1] == [vlm_image] for call in encode_calls)
assert len(pipeline.transformer.latent_inputs) == expected_transformer_calls
for item in pipeline.transformer.latent_inputs:
expected_prefix = condition.clean_latent.expand(item.shape[0], -1, -1, -1, -1)
assert torch.equal(item[:, :, :1], expected_prefix)
assert torch.equal(latents[:, :, :1], condition.clean_latent)
assert torch.equal(latents[:, :, 1:], torch.full((1, 1, 2, 1, 1), 14.0))
def test_t2i_decodes_and_returns_the_unique_frame():
from vllm_omni.diffusion.models.lingbot_video import LingBotGenerationMode
pipeline = _make_pipeline()
pipeline.transformer = _RecordingTransformer()
pipeline.scheduler = _CorruptingScheduler()
pipeline.progress_bar = lambda timesteps: timesteps
pipeline.prepare_latents = lambda *args, **kwargs: torch.zeros(1, 1, 1, 2, 2)
pipeline.encode_prompt = lambda *args, **kwargs: (
torch.ones(1, 2, 4),
torch.ones(1, 2, dtype=torch.long),
)
pipeline._decode_latents = lambda latents: torch.ones(1, 16, 16, 3)
image = pipeline._generate(
prompt="a still robot",
mode=LingBotGenerationMode.T2I,
height=16,
width=16,
num_frames=1,
num_inference_steps=2,
guidance_scale=1.0,
shift=3.0,
output_type="pt",
)
assert image.shape == (16, 16, 3)
def test_forward_rejects_multi_request_batches():
pipeline = _make_pipeline()
first = OmniDiffusionRequest(
@@ -527,45 +211,6 @@ def test_forward_rejects_multi_request_batches():
pipeline.forward(DiffusionRequestBatch([first, second]))
def test_forward_marks_invalid_request_contract_as_client_error():
from vllm_omni.errors import OmniClientError
pipeline = _make_pipeline()
request = _make_request_batch(
{
"prompt": "unsupported image edit",
"modalities": ["image"],
"multi_modal_data": {"image": Image.new("RGB", (32, 32), color="blue")},
},
height=192,
width=320,
num_frames=1,
)
with pytest.raises(OmniClientError, match="does not accept") as exc_info:
pipeline.forward(request)
assert exc_info.value.status_code == 400
def test_forward_marks_unsupported_output_count_as_client_error():
from vllm_omni.errors import OmniClientError
pipeline = _make_pipeline()
request = _make_request_batch(
{"prompt": "motion", "modalities": ["video"]},
height=192,
width=320,
num_frames=9,
num_outputs_per_prompt=2,
)
with pytest.raises(OmniClientError, match="one output per prompt") as exc_info:
pipeline.forward(request)
assert exc_info.value.status_code == 400
def test_load_weights_rejects_external_weight_stream():
pipeline = _make_pipeline()
@@ -1,350 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import pytest
from PIL import Image
from vllm_omni.diffusion.models.lingbot_video.request_utils import (
LingBotGenerationMode,
caption_from_lingbot_prompt,
normalize_lingbot_num_frames,
normalize_lingbot_request,
resolve_lingbot_mode,
resolve_lingbot_num_frames,
resolve_lingbot_output_dimensions,
resolve_lingbot_size,
)
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def _request(prompt, **sampling_overrides):
return SimpleNamespace(
prompt=prompt,
sampling_params=OmniDiffusionSamplingParams(**sampling_overrides),
)
def _normalize(prompt, **sampling_overrides):
return normalize_lingbot_request(
_request(prompt, **sampling_overrides),
default_negative_prompt="video default",
default_image_negative_prompt="image default",
)
@pytest.mark.parametrize(
("prompt", "expected"),
[
({"prompt": "still", "modalities": ["image"]}, LingBotGenerationMode.T2I),
({"prompt": "motion", "modalities": ["video"]}, LingBotGenerationMode.T2V),
(
{
"prompt": "motion",
"modalities": ["video"],
"multi_modal_data": {"image": Image.new("RGB", (8, 8))},
},
LingBotGenerationMode.TI2V,
),
],
)
def test_resolve_lingbot_mode(prompt, expected):
assert resolve_lingbot_mode(prompt) is expected
@pytest.mark.parametrize(
("prompt", "message"),
[
(
{"prompt": "x", "modalities": ["image"], "multi_modal_data": {"image": Image.new("RGB", (8, 8))}},
"does not accept",
),
(
{
"prompt": "x",
"modalities": ["video"],
"multi_modal_data": {"image": [Image.new("RGB", (8, 8)), Image.new("RGB", (8, 8))]},
},
"exactly one",
),
({"prompt": "x", "modalities": ["video"], "multi_modal_data": {"video": [1]}}, "reference modalities"),
({"prompt": "x", "multi_modal_data": {"image": Image.new("RGB", (8, 8))}}, "explicitly set"),
({"prompt": "x", "modalities": ["audio"]}, "output modality"),
],
)
def test_resolve_lingbot_mode_rejects_invalid_media_combinations(prompt, message):
with pytest.raises(ValueError, match=message):
resolve_lingbot_mode(prompt)
def test_caption_from_lingbot_prompt_preserves_text_and_compact_json():
assert caption_from_lingbot_prompt("plain text") == "plain text"
assert caption_from_lingbot_prompt({"caption": "caption only", "duration": 4}) == "caption only"
assert caption_from_lingbot_prompt({"caption": {"中文": True, "nested": {"value": 2}}}) == (
'{"中文":true,"nested":{"value":2}}'
)
assert (
caption_from_lingbot_prompt(
{"scene": "海边", "camera": {"move": "pan"}, "duration": 5, "fps": 24, "size": "320x192"}
)
== '{"scene":"海边","camera":{"move":"pan"}}'
)
@pytest.mark.parametrize("prompt", [{"duration": 4, "fps": 24}, {"caption": ""}, ""])
def test_caption_from_lingbot_prompt_rejects_empty_caption(prompt):
with pytest.raises(ValueError, match="empty|must not be empty"):
caption_from_lingbot_prompt(prompt)
@pytest.mark.parametrize(
("num_frames", "expected"),
[(1, 1), (2, 5), (4, 5), (5, 5), (96, 97), (98, 101), (121, 121)],
)
def test_normalize_lingbot_num_frames_rounds_up_to_4n_plus_1(num_frames, expected):
assert normalize_lingbot_num_frames(num_frames) == expected
@pytest.mark.parametrize("num_frames", [0, -1, 1.5, "invalid"])
def test_normalize_lingbot_num_frames_rejects_invalid_values(num_frames):
with pytest.raises(ValueError, match="num_frames"):
normalize_lingbot_num_frames(num_frames)
@pytest.mark.parametrize(
("duration", "fps", "expected"),
[(4, 24, 97), (5, 24, 121), (8, 24, 193), (4, 12, 49), (4, 30, 121)],
)
def test_resolve_lingbot_num_frames_rounds_up(duration, fps, expected):
assert resolve_lingbot_num_frames(duration, fps) == expected
@pytest.mark.parametrize(("duration", "fps"), [(0, 24), (-1, 24), (0.01, 24), (1, 0), (1, -1)])
def test_resolve_lingbot_num_frames_rejects_non_positive_or_empty_results(duration, fps):
with pytest.raises(ValueError, match="duration|fps|at least one"):
resolve_lingbot_num_frames(duration, fps)
@pytest.mark.parametrize(
("resolution", "ratio", "expected"),
[
("192p", "1:1", (192, 192)),
("192p", "9:16", (192, 320)),
("192p", "16:9", (320, 192)),
("192p", "3:4", (192, 256)),
("192p", "4:3", (256, 192)),
("480p", "1:1", (480, 480)),
("480p", "9:16", (480, 832)),
("480p", "16:9", (832, 480)),
("480p", "3:4", (480, 640)),
("480p", "4:3", (640, 480)),
("720p", "1:1", (736, 736)),
("720p", "9:16", (736, 1280)),
("720p", "16:9", (1280, 736)),
("720p", "3:4", (736, 960)),
("720p", "4:3", (960, 736)),
("1080p", "1:1", (1088, 1088)),
("1080p", "9:16", (1088, 1920)),
("1080p", "16:9", (1920, 1088)),
("1080p", "3:4", (1088, 1440)),
("1080p", "4:3", (1440, 1088)),
("2k", "1:1", (1440, 1440)),
("2k", "9:16", (1440, 2560)),
("2k", "16:9", (2560, 1440)),
("2k", "3:4", (1440, 1920)),
("2k", "4:3", (1920, 1440)),
("4k", "1:1", (2176, 2176)),
("4k", "9:16", (2176, 3840)),
("4k", "16:9", (3840, 2176)),
("4k", "3:4", (2176, 2880)),
("4k", "4:3", (2880, 2176)),
],
)
def test_resolve_lingbot_size_supports_all_official_presets(resolution, ratio, expected):
assert resolve_lingbot_size(resolution=resolution, ratio=ratio) == expected
def test_resolve_lingbot_size_validates_sources_and_alignment():
assert resolve_lingbot_size(size="320x192") == (192, 320)
assert resolve_lingbot_size(width=320, height=192) == (192, 320)
with pytest.raises(ValueError, match="provided together"):
resolve_lingbot_size(width=320)
with pytest.raises(ValueError, match="provided together"):
resolve_lingbot_size(resolution="192p")
with pytest.raises(ValueError, match="mutually exclusive"):
resolve_lingbot_size(width=320, height=192, resolution="192p", ratio="9:16")
with pytest.raises(ValueError, match="multiples of 16"):
resolve_lingbot_size(width=321, height=192)
def test_resolve_lingbot_output_dimensions_matches_request_priority():
assert resolve_lingbot_output_dimensions(
sampling_width=480,
sampling_height=480,
prompt_fields={"resolution": "720p", "ratio": "16:9"},
) == (736, 1280)
assert resolve_lingbot_output_dimensions(
sampling_width=480,
sampling_height=480,
prompt_fields={"width": 320, "height": 192},
) == (320, 192)
assert resolve_lingbot_output_dimensions(
sampling_width=320,
sampling_height=192,
) == (320, 192)
assert resolve_lingbot_output_dimensions() == (480, 480)
def test_normalize_lingbot_request_uses_mode_defaults_and_preserves_empty_negative():
t2i = _normalize(
{"prompt": "still", "modalities": ["image"], "negative_prompt": ""},
width=320,
height=192,
num_frames=1,
)
assert t2i.mode is LingBotGenerationMode.T2I
assert t2i.num_frames == 1
assert t2i.negative_prompt == ""
t2v = _normalize(
{"prompt": "motion", "modalities": ["video"]},
width=320,
height=192,
num_frames=9,
)
assert t2v.mode is LingBotGenerationMode.T2V
assert t2v.negative_prompt == "video default"
def test_normalize_lingbot_request_rejects_explicit_zero_flow_shift():
with pytest.raises(ValueError, match="shift.*positive"):
_normalize(
{"prompt": "still", "modalities": ["image"]},
width=320,
height=192,
num_frames=1,
extra_args={"flow_shift": 0},
)
def test_normalize_lingbot_request_rejects_explicit_zero_width():
with pytest.raises(ValueError, match="width"):
_normalize(
{"prompt": "still", "modalities": ["image"]},
width=0,
height=192,
num_frames=1,
)
def test_normalize_lingbot_request_rounds_all_video_frame_counts():
config = _normalize(
{"prompt": "motion", "modalities": ["video"]},
width=320,
height=192,
num_frames=96,
fps=24,
frame_rate=24,
)
assert config.num_frames == 97
config = _normalize(
{"prompt": "motion", "modalities": ["video"]},
width=320,
height=192,
num_frames=121,
fps=24,
frame_rate=24,
)
assert config.num_frames == 121
def test_normalize_lingbot_request_prefers_explicit_frames_over_duration():
online = _normalize(
{"prompt": "motion", "modalities": ["video"]},
width=320,
height=192,
num_frames=9,
fps=24,
extra_args={"duration": 5},
)
assert online.num_frames == 9
offline = _normalize(
{
"prompt": {"caption": "motion", "duration": 5, "num_frames": 9},
"modalities": ["video"],
},
width=320,
height=192,
fps=24,
)
assert offline.num_frames == 9
def test_normalize_lingbot_request_uses_duration_then_model_default():
duration = _normalize(
{"prompt": {"caption": "motion", "duration": 5}, "modalities": ["video"]},
width=320,
height=192,
fps=24,
)
assert duration.num_frames == 121
default = _normalize(
{"prompt": "motion", "modalities": ["video"]},
width=320,
height=192,
)
assert default.num_frames == 81
def test_normalize_lingbot_request_rejects_contract_conflicts():
with pytest.raises(ValueError, match="text-to-image.*duration"):
_normalize(
{"prompt": {"caption": "still", "duration": 4}, "modalities": ["image"]},
width=320,
height=192,
num_frames=1,
)
with pytest.raises(ValueError, match="text-to-image requires"):
_normalize(
{"prompt": "still", "modalities": ["image"]},
width=320,
height=192,
num_frames=5,
)
with pytest.raises(ValueError, match="seconds.*duration"):
_normalize(
{"prompt": "motion", "modalities": ["video"]},
width=320,
height=192,
num_frames=96,
extra_args={"seconds": 4, "duration": 4},
)
def test_normalize_lingbot_request_model_size_overrides_sampling_dimensions():
# A prompt-provided size overrides the sampling width/height.
config = _normalize(
{"prompt": {"caption": "motion", "size": "320x192"}, "modalities": ["video"]},
width=480,
height=480,
num_frames=81,
)
assert (config.height, config.width) == (192, 320)
# Without a prompt size, the sampling dimensions are used.
config = _normalize(
{"prompt": "motion", "modalities": ["video"]},
width=480,
height=480,
num_frames=81,
)
assert (config.height, config.width) == (480, 480)
@@ -1,266 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import json
import numpy as np
import pytest
import torch
from PIL import Image
from vllm_omni.diffusion.models.longcat_video.longcat_video_avatar_transformer import (
LongCatVideoAvatarTransformer3DModel,
_read_config,
replace_linear_with_quantized,
)
from vllm_omni.diffusion.models.longcat_video.pipeline_longcat_video_avatar import (
_avatar_model_allow_patterns,
_build_multi_speaker_ref_target_masks,
_default_at2v_shape,
_infer_asset_root_from_path,
_prepare_multi_speaker_audio_arrays,
_resolve_num_segments,
prepare_longcat_video_avatar_model_for_omni,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@pytest.mark.parametrize(
("resolution", "expected_shape"),
[
("480p", (480, 832)),
("720p", (736, 1248)),
],
)
def test_longcat_video_avatar_at2v_default_shape_uses_resolution_bucket(
resolution: str,
expected_shape: tuple[int, int],
):
assert _default_at2v_shape(resolution) == expected_shape
@pytest.mark.parametrize(
("use_int8", "expected_weight_dir", "unexpected_weight_dir"),
[
(True, "base_model_int8/*", "base_model/*"),
(False, "base_model/*", "base_model_int8/*"),
],
)
def test_longcat_video_avatar_allow_patterns_download_one_weight_set(
use_int8: bool,
expected_weight_dir: str,
unexpected_weight_dir: str,
):
allow_patterns = _avatar_model_allow_patterns(use_int8)
assert expected_weight_dir in allow_patterns
assert unexpected_weight_dir not in allow_patterns
assert "whisper-large-v3/*" not in allow_patterns
assert "whisper-large-v3/model.safetensors" in allow_patterns
assert "whisper-large-v3/pytorch_model.bin" not in allow_patterns
assert "whisper-large-v3/flax_model.msgpack" not in allow_patterns
assert "vocal_separator/*" in allow_patterns
def test_longcat_video_avatar_prepare_model_adds_omni_metadata(tmp_path):
model_dir = tmp_path / "LongCat-Video-Avatar-1.5"
base_model_dir = model_dir / "base_model_int8"
base_model_dir.mkdir(parents=True)
(base_model_dir / "config.json").write_text('{"model_type": "longcat_avatar"}', encoding="utf-8")
(model_dir / "model_index.json").write_text("{}", encoding="utf-8")
prepared = prepare_longcat_video_avatar_model_for_omni(str(model_dir), use_int8=True)
assert prepared == str(model_dir)
model_index = json.loads((model_dir / "model_index.json").read_text(encoding="utf-8"))
assert model_index["_class_name"] == "LongCatVideoAvatarPipeline"
assert model_index["_diffusers_version"] == "0.38.0"
transformer_config = json.loads((model_dir / "transformer" / "config.json").read_text(encoding="utf-8"))
assert transformer_config == {"model_type": "longcat_avatar"}
def test_longcat_video_avatar_infers_repo_root_from_official_asset_path(tmp_path):
input_json = tmp_path / "assets" / "avatar" / "single_example_1.json"
input_json.parent.mkdir(parents=True)
input_json.write_text("{}", encoding="utf-8")
assert _infer_asset_root_from_path(input_json) == tmp_path
def test_longcat_video_avatar_falls_back_to_json_parent_without_official_layout(tmp_path):
input_json = tmp_path / "single_example_1.json"
input_json.write_text("{}", encoding="utf-8")
assert _infer_asset_root_from_path(input_json) == tmp_path
@pytest.mark.parametrize(
("value", "audio_duration", "expected"),
[
(None, 26.6, 1),
("1", 26.6, 1),
(3, 26.6, 3),
("auto", 3.0, 1),
("auto", 23.8, 8),
("auto", 26.6, 9),
],
)
def test_longcat_video_avatar_resolve_num_segments(
value,
audio_duration: float,
expected: int,
):
assert (
_resolve_num_segments(
value,
audio_duration=audio_duration,
num_frames=93,
num_cond_frames=13,
save_fps=25,
)
== expected
)
def test_longcat_video_avatar_read_config_filters_non_constructor_metadata(tmp_path):
config_path = tmp_path / "config.json"
config_path.write_text(
json.dumps(
{
"_class_name": "LongCatVideoAvatarTransformer3DModel",
"architectures": ["LongCatVideoAvatarTransformer3DModel"],
"_diffusers_version": "0.35.1",
"model_max_length": 512,
"hidden_size": 4096,
"depth": 48,
"num_heads": 32,
}
),
encoding="utf-8",
)
config = _read_config(config_path)
assert config == {
"hidden_size": 4096,
"depth": 48,
"num_heads": 32,
}
def test_longcat_video_avatar_load_weights_supports_int8_buffers():
model = LongCatVideoAvatarTransformer3DModel(
hidden_size=4,
depth=0,
num_heads=1,
caption_channels=4,
intermediate_dim=4,
output_dim=4,
audio_channel=4,
context_tokens=1,
)
replace_linear_with_quantized(model)
buffer_name = "t_embedder.mlp.0.weight_int8"
buffers = dict(model.named_buffers())
loaded_weight = torch.ones_like(buffers[buffer_name])
loaded_params = model.load_weights([(buffer_name, loaded_weight)])
assert buffer_name in loaded_params
assert torch.equal(dict(model.named_buffers())[buffer_name], loaded_weight)
def test_longcat_video_avatar_multi_speaker_para_audio_arrays_are_aligned():
left = np.ones(3, dtype=np.float32)
right = np.ones(5, dtype=np.float32) * 2
left_out, right_out = _prepare_multi_speaker_audio_arrays(
[left, right],
audio_type="para",
generate_duration=0.5,
sample_rate=10,
)
assert left_out.tolist() == [1, 1, 1, 0, 0]
assert right_out.tolist() == [2, 2, 2, 2, 2]
def test_longcat_video_avatar_multi_speaker_add_audio_arrays_are_sequential():
left = np.ones(3, dtype=np.float32)
right = np.ones(2, dtype=np.float32) * 2
left_out, right_out = _prepare_multi_speaker_audio_arrays(
[left, right],
audio_type="add",
generate_duration=0.1,
sample_rate=10,
)
assert left_out.tolist() == [1, 1, 1, 0, 0]
assert right_out.tolist() == [0, 0, 0, 2, 2]
def test_longcat_video_avatar_multi_speaker_masks_use_explicit_bboxes():
image = Image.new("RGB", (8, 6))
masks, use_background = _build_multi_speaker_ref_target_masks(
image,
{
"person1": [1, 1, 4, 3],
"person2": [2, 4, 5, 7],
},
)
assert not use_background
assert masks.shape == (3, 6, 8)
assert masks[0, 1:4, 1:3].sum() == 6
assert masks[1, 2:5, 4:7].sum() == 9
assert masks[2, 0, 0] == 1
assert masks[2, 2, 2] == 0
def test_longcat_video_avatar_multi_speaker_masks_default_to_left_right_split():
image = Image.new("RGB", (10, 8))
masks, use_background = _build_multi_speaker_ref_target_masks(image, None)
assert not use_background
assert masks.shape == (3, 8, 10)
assert masks[0, :, :5].sum() > 0
assert masks[0, :, 5:].sum() == 0
assert masks[1, :, :5].sum() == 0
assert masks[1, :, 5:].sum() > 0
def test_longcat_video_avatar_transformer_accepts_multi_speaker_masks():
model = LongCatVideoAvatarTransformer3DModel(
hidden_size=4,
depth=0,
num_heads=1,
caption_channels=4,
intermediate_dim=4,
output_dim=4,
audio_channel=4,
context_tokens=1,
)
hidden_states = torch.randn(1, 16, 5, 4, 4)
timestep = torch.zeros(1, 5)
encoder_hidden_states = torch.randn(1, 1, 4, 4)
audio_embs = torch.randn(2, 5, 5, 12, 4)
ref_target_masks = torch.zeros(3, 4, 4)
ref_target_masks[0, :, :2] = 1
ref_target_masks[1, :, 2:] = 1
ref_target_masks[2] = torch.where(ref_target_masks[0] + ref_target_masks[1] > 0, 0, 1)
output = model(
hidden_states=hidden_states,
timestep=timestep,
encoder_hidden_states=encoder_hidden_states,
num_cond_latents=1,
audio_embs=audio_embs,
ref_target_masks=ref_target_masks,
)
assert output.shape == hidden_states.shape
@@ -484,182 +484,6 @@ def test_shifted_sigma_schedule_matches_reference_values():
)
def test_base_schedule_overrides_the_uniform_sigma_positions():
from vllm_omni.diffusion.models.minimax_h3.time_request import (
minimax_h3_time_shift_sigmas,
)
base_schedule = [1.0, 0.7, 0.4, 0.15, 0.0]
video = minimax_h3_time_shift_sigmas(
num_steps=len(base_schedule),
shift_scale=12.0,
base_schedule=base_schedule,
)
audio = minimax_h3_time_shift_sigmas(
num_steps=len(base_schedule),
shift_scale=3.0,
base_schedule=base_schedule,
)
assert video == pytest.approx([1.0, 0.9655172, 0.8888889, 0.6792453, 0.0], abs=1e-6)
assert audio == pytest.approx([1.0, 0.875, 0.6666667, 0.3461539, 0.0], abs=1e-6)
assert video != pytest.approx(
minimax_h3_time_shift_sigmas(num_steps=len(base_schedule), shift_scale=12.0),
abs=1e-6,
)
@pytest.mark.parametrize(
"base_schedule",
[
[0.9, 0.5, 0.0],
[1.0, 0.5, 0.1],
[1.0, 0.5, 0.5, 0.0],
[1.0],
[],
[1.0, float("nan"), 0.4, 0.0],
[1.0, float("inf"), 0.4, 0.0],
],
)
def test_base_schedule_rejects_malformed_positions(base_schedule):
from vllm_omni.diffusion.models.minimax_h3.time_request import (
minimax_h3_time_shift_sigmas,
)
with pytest.raises(ValueError):
minimax_h3_time_shift_sigmas(
num_steps=max(len(base_schedule), 1),
shift_scale=12.0,
base_schedule=base_schedule,
)
def _distilled_pipeline(diffuse_calls, base_schedule_by_partition):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.sched import DMD2SigmaSchedule
schedules = {
partition: None if positions is None else DMD2SigmaSchedule.from_positions(positions)
for partition, positions in base_schedule_by_partition.items()
}
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "combined"
pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
pipeline.default_video_shift = 12.0
pipeline.default_audio_shift = 3.0
pipeline.device = torch.device("cpu")
pipeline.od_config = SimpleNamespace()
pipeline._base_schedule_by_partition = schedules
pipeline._quality_policy = Mock()
pipeline._quality_policy.resolve.return_value = SimpleNamespace(cache_dit=None)
pipeline._cache_dit_runtime = SimpleNamespace(prepare=lambda spec: None)
pipeline.encode_prompt = Mock(
return_value=(
torch.ones(1, 2),
torch.ones(1, dtype=torch.long),
)
)
def diffuse(**kwargs):
diffuse_calls.append(kwargs)
return torch.zeros(1), torch.zeros(1)
pipeline.diffuse = diffuse
pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
return pipeline
def _t2va_batch(num_inference_steps=None):
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
sampling = OmniDiffusionSamplingParams(
quality="lossless",
width=1344,
height=768,
fps=24,
num_frames=124,
num_inference_steps=num_inference_steps,
extra_args={"task": "t2va", "aspect_ratio": "16:9"},
)
return DiffusionRequestBatch(
[
OmniDiffusionRequest(
prompt="distilled schedule",
sampling_params=sampling,
request_id="distilled",
)
]
)
def test_base_schedule_is_scoped_to_the_serving_partition():
"""A distilled FL2VA must not drag a regular Ref2VA onto its step count."""
distilled = [1.0, 0.7, 0.4, 0.15, 0.0]
pipeline = _distilled_pipeline([], {"fl2va": distilled, "ref2va": None})
assert pipeline._base_schedule_for_task("t2va").base_schedule == tuple(distilled)
assert pipeline._base_schedule_for_task("fl2va").base_schedule == tuple(distilled)
assert pipeline._base_schedule_for_task("ref2va") is None
def test_partially_constructed_pipeline_falls_back_to_the_uniform_schedule():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
assert pipeline._base_schedule_for_task("t2va") is None
def test_distilled_forward_reports_denoising_steps_not_sigma_boundaries():
diffuse_calls = []
base_schedule = [1.0, 0.7, 0.4, 0.15, 0.0]
pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": base_schedule, "ref2va": None})
pipeline.forward(_t2va_batch())
assert diffuse_calls[0]["base_schedule"] == tuple(base_schedule)
# Five boundaries describe four denoising steps.
assert diffuse_calls[0]["num_steps"] == 4
pipeline._quality_policy.resolve.assert_called_once_with(
quality="lossless",
num_inference_steps=4,
extra_args={"task": "t2va", "aspect_ratio": "16:9"},
)
def test_distilled_forward_accepts_the_matching_explicit_step_count():
diffuse_calls = []
pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": [1.0, 0.7, 0.4, 0.15, 0.0], "ref2va": None})
pipeline.forward(_t2va_batch(num_inference_steps=4))
assert diffuse_calls[0]["num_steps"] == 4
def test_distilled_forward_rejects_a_mismatched_explicit_step_count():
from vllm_omni.errors import OmniClientError
pipeline = _distilled_pipeline([], {"fl2va": [1.0, 0.7, 0.4, 0.15, 0.0], "ref2va": None})
with pytest.raises(OmniClientError, match="must be 4 or omitted"):
pipeline.forward(_t2va_batch(num_inference_steps=50))
def test_absent_base_schedule_key_differs_from_an_empty_list():
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_read_base_schedule,
)
assert _read_base_schedule({}) is None
assert _read_base_schedule({"base_schedule": [1.0, 0.5, 0.0]}).base_schedule == (1.0, 0.5, 0.0)
with pytest.raises(ValueError):
_read_base_schedule({"base_schedule": []})
def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
MiniMaxH3Attention,
@@ -668,10 +492,6 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
class FakeBackend:
supports_prefix_kv_slicing = True
@classmethod
def supports_packed_mask_free(cls) -> bool:
return False
class FakeAttention(torch.nn.Module):
attn_backend = FakeBackend
@@ -706,48 +526,6 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
assert attention.attention.metadata.extra["valid_kv_length"] == 5
def test_packed_attention_skips_mask_for_packed_mask_free_backend():
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
MiniMaxH3Attention,
)
class FakeBackend:
supports_prefix_kv_slicing = False
@classmethod
def supports_packed_mask_free(cls) -> bool:
return True
class FakeAttention(torch.nn.Module):
attn_backend = FakeBackend
def __init__(self):
super().__init__()
self.metadata = None
def forward(self, query, key, value, metadata):
self.metadata = metadata
return query
attention = object.__new__(MiniMaxH3Attention)
torch.nn.Module.__init__(attention)
attention.attention = FakeAttention()
q = torch.randn(8, 2, 4)
attention._run_packed_attention(
q,
q,
q,
cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
max_seqlen=5,
packed_total=8,
)
assert attention.attention.metadata.attn_mask is None
assert attention.attention.metadata.extra["valid_kv_length"] == 5
assert attention.attention.metadata.extra["npu_attn_varlen"] is True
def test_packed_attention_keeps_padding_mask_for_other_backends():
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
MiniMaxH3Attention,
@@ -756,10 +534,6 @@ def test_packed_attention_keeps_padding_mask_for_other_backends():
class FakeBackend:
supports_prefix_kv_slicing = False
@classmethod
def supports_packed_mask_free(cls) -> bool:
return False
class FakeAttention(torch.nn.Module):
attn_backend = FakeBackend
@@ -1422,38 +1196,8 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti
)
assert prepared[0]["duration_seconds"] == pytest.approx(expected_duration)
assert prepared[0]["audio_duration_seconds"] == pytest.approx(min(expected_duration, 209 / 24.0))
assert prepared[0]["start_time_seconds"] == pytest.approx(start_time or 0.0)
assert transcode_calls[0][1]["duration_seconds"] == pytest.approx(expected_duration)
assert transcode_calls[0][1]["target_frame_count"] == 209
def test_ref2va_qwen_sampling_decodes_prepared_video_once(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
decoded = np.arange(25 * 2 * 2 * 3, dtype=np.uint8).reshape(25, 2, 2, 3)
load_calls = []
monkeypatch.setattr(
reference_video_module,
"_probe_video",
lambda _path: {"frame_count": len(decoded)},
)
monkeypatch.setattr(
reference_video_module,
"load_video_frames",
lambda path: load_calls.append(path) or decoded,
)
sampled = reference_video_module.sample_reference_video_frames("prepared.mp4")
assert load_calls == ["prepared.mp4"]
for frame, expected in zip(
sampled["frames"],
(decoded[0], decoded[12], decoded[24]),
strict=True,
):
np.testing.assert_array_equal(frame, expected)
assert sampled["block_timestamps"] == [0.25, 1.0]
def test_ref2va_two_video_recipe_tolerates_container_rounding(monkeypatch, tmp_path):
@@ -1693,47 +1437,6 @@ def test_g4_standalone_audio_duration_and_total_duration_contract():
)
def test_ref2va_audio_duration_validation_precedes_rank_branch(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
pipeline = object.__new__(pipeline_module.MiniMaxH3Pipeline)
pipeline.device = torch.device("cpu")
monkeypatch.setattr(pipeline_module, "_dit_rank_world", lambda: (None, 1, 2))
waveform = torch.zeros(1, 10)
with pytest.raises(ValueError, match="max_duration_seconds must be positive"):
pipeline._encode_audio_conditions(
[(waveform, 10)],
max_duration_seconds=0,
)
def test_ref2va_standalone_audio_condition_is_bounded_to_output_duration():
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
pipeline.device = torch.device("cpu")
observed = []
class FakeAudioVAE:
def encode_waveform(self, waveform, sample_rate):
observed.append((waveform.clone(), sample_rate))
return waveform, waveform.shape[-1]
pipeline.audio_vae = FakeAudioVAE()
waveform = torch.arange(40, dtype=torch.float32).reshape(1, 40)
encoded, lengths = pipeline._encode_audio_conditions(
[(waveform, 10)],
max_duration_seconds=2.5,
)
assert observed[0][0].shape == (1, 25)
assert observed[0][1] == 10
assert encoded.shape == (1, 25)
assert lengths == [25]
@pytest.mark.parametrize(
("field", "value", "message"),
[
@@ -19,7 +19,7 @@ MODEL_ENV = "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL"
def _assert_joint_output(outputs):
assert len(outputs) == 1
frames = np.asarray(outputs[0].images[0])
assert frames.shape == (107, 256, 448, 3)
assert frames.shape == (39, 256, 448, 3)
multimodal = outputs[0].multimodal_output
assert multimodal is not None
assert np.asarray(multimodal["audio"]).shape[1] == 2
@@ -122,58 +122,6 @@ def test_minimax_h3_t2va_fp8_single_gpu_smoke():
output_type="np",
extra_args={
"task": "t2va",
"duration": 4.0,
"aspect_ratio": "16:9",
"flow_shift": 12.0,
"audio_flow_shift": 3.0,
},
),
use_tqdm=False,
)
finally:
engine.close()
_assert_joint_output(outputs)
@pytest.mark.skipif(
not os.environ.get(MODEL_ENV),
reason=f"set {MODEL_ENV} to an authorized FL2VA checkpoint path",
)
def test_minimax_h3_t2va_teacache_single_gpu_smoke():
import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
if torch.accelerator.device_count() < 1:
pytest.skip("MiniMax H3 TeaCache smoke requires a CUDA device")
engine = Omni(
model=os.environ[MODEL_ENV],
cache_backend="tea_cache",
cache_config={"rel_l1_thresh": 0.17},
parallel_config=DiffusionParallelConfig(ulysses_degree=1),
trust_remote_code=True,
enable_cpu_offload=True,
enforce_eager=True,
)
try:
outputs = engine.generate(
"A quiet cinematic night scene with matching ambient sound.",
OmniDiffusionSamplingParams(
height=256,
width=448,
num_frames=29,
fps=24,
num_inference_steps=2,
seed=42,
output_type="np",
extra_args={
"task": "t2va",
"duration": 4.0,
"aspect_ratio": "16:9",
"flow_shift": 12.0,
"audio_flow_shift": 3.0,
},
@@ -69,7 +69,6 @@ def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
attention = object.__new__(MiniMaxH3Attention)
nn.Module.__init__(attention)
attention.rot_dim = 96
attention.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
attention.rope._forward_method = attention.rope.forward_native
@@ -1,52 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm_omni.diffusion.sched import DMD2SigmaSchedule
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_step_count_is_intervals_not_boundaries():
schedule = DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])
assert schedule.num_inference_steps == 4
@pytest.mark.parametrize(
"positions",
[
[],
[1.0],
[0.9, 0.5, 0.0],
[1.0, 0.5, 0.1],
[1.0, 0.4, 0.6, 0.0],
[1.0, 0.5, 0.5, 0.0],
[1.0, float("nan"), 0.4, 0.0],
[1.0, float("inf"), 0.4, 0.0],
],
)
def test_malformed_positions_are_rejected(positions):
with pytest.raises(ValueError):
DMD2SigmaSchedule.from_positions(positions)
def test_one_base_schedule_drives_several_shift_scales():
schedule = DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])
assert schedule.shifted_sigmas(12.0) == pytest.approx([1.0, 0.9655172, 0.8888889, 0.6792453, 0.0], abs=1e-6)
assert schedule.shifted_sigmas(3.0) == pytest.approx([1.0, 0.875, 0.6666667, 0.3461539, 0.0], abs=1e-6)
def test_non_positive_shift_scale_is_rejected():
schedule = DMD2SigmaSchedule.from_positions([1.0, 0.5, 0.0])
with pytest.raises(ValueError):
schedule.shifted_sigmas(0.0)
def test_absent_metadata_key_means_undistilled_but_empty_is_malformed():
assert DMD2SigmaSchedule.from_metadata({}) is None
assert DMD2SigmaSchedule.from_metadata({"base_schedule": [1.0, 0.5, 0.0]}).base_schedule == (1.0, 0.5, 0.0)
with pytest.raises(ValueError):
DMD2SigmaSchedule.from_metadata({"base_schedule": []})
@@ -15,7 +15,6 @@ from vllm_omni.diffusion.data import (
DiffusionParallelConfig,
OmniDiffusionConfig,
)
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.model_metadata import (
HUNYUAN_IMAGE3_MAX_INPUT_IMAGES,
QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
@@ -125,11 +124,6 @@ class TestCreateDefaultDiffusion:
assert ea["enforce_eager"] is True
assert ea["lora_path"] == "/tmp/lora"
def test_diffusion_kv_mode_roundtrip(self):
od = _roundtrip_diffusion_config(model="x", diffusion_kv_mode="paged_scheduler")
assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline")
+1 -74
View File
@@ -25,7 +25,6 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.interface import (
CachedRequestData,
NewRequestData,
_AdmissionWaitDecision,
)
from vllm_omni.diffusion.sched.interface import (
DiffusionSchedulerOutput as RealDiffusionSchedulerOutput,
@@ -76,34 +75,6 @@ class MockScheduler:
def num_running_requests(self) -> int:
return 0
def get_admission_wait_decision(
self,
*,
now: float,
dp_concurrent: bool = False,
) -> _AdmissionWaitDecision:
del dp_concurrent
return _AdmissionWaitDecision(
should_wait=True,
deadline=now + 0.5,
stable_window_s=0.05,
max_batch=self.max_num_running_reqs,
)
def should_end_admission_wait(
self,
decision: _AdmissionWaitDecision,
*,
now: float,
stable_since: float,
) -> bool:
waiting = self.num_waiting_requests()
return (
waiting >= decision.max_batch
or (waiting > 0 and now - stable_since >= decision.stable_window_s)
or (decision.deadline is not None and now >= decision.deadline)
)
def schedule(self) -> DiffusionSchedulerOutput:
if not self._waiting_queue:
return DiffusionSchedulerOutput(step_id=self._step_id)
@@ -137,17 +108,6 @@ class MockScheduler:
pass
def _make_admission_engine(pre_process_func) -> DiffusionEngine:
engine = object.__new__(DiffusionEngine)
engine.scheduler = MockScheduler()
engine.pre_process_func = pre_process_func
engine._out_streams = {}
engine._rpc_lock = threading.RLock()
engine._cv = threading.Condition(engine._rpc_lock)
engine._closed = False
return engine
class _BatchCapablePipeline:
supports_request_batch = True
@@ -565,11 +525,6 @@ class TestRequestBatchAdmission:
with pytest.raises(ValueError, match="request_batch_max_wait_ms"):
OmniDiffusionConfig(model="test", request_batch_max_wait_ms=-1.0)
@pytest.mark.parametrize("bad_wait", [float("nan"), float("inf"), float("-inf")])
def test_config_rejects_nonfinite_request_batch_max_wait_ms(self, bad_wait: float) -> None:
with pytest.raises(ValueError, match="request_batch_max_wait_ms"):
OmniDiffusionConfig(model="test", request_batch_max_wait_ms=bad_wait)
def test_config_normalizes_request_batch_max_wait_ms_to_float(self) -> None:
config = OmniDiffusionConfig(model="test", request_batch_max_wait_ms=5)
@@ -637,7 +592,7 @@ class TestRequestBatchAdmission:
start = time.monotonic()
with engine._cv:
engine._wait_for_admission_if_needed_locked()
engine._wait_for_request_batch_admission_locked()
waited_s = time.monotonic() - start
# Stable-window exit (~50ms), not the full 1000ms deadline.
@@ -728,34 +683,6 @@ async def _consume_final_output(generator):
return final_output
@pytest.mark.cpu
@pytest.mark.parametrize("entrypoint", ["add_request", "async_add_req_and_stream_response"])
def test_engine_admission_preprocesses_request_once(entrypoint: str) -> None:
raw_request = OmniDiffusionRequest(
prompt="raw",
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
request_id="admission",
)
prepared_request = OmniDiffusionRequest(
prompt="prepared",
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
request_id="admission",
prepared_layout=object(),
)
preprocess_calls = []
def preprocess(request):
preprocess_calls.append(request)
return prepared_request
engine = _make_admission_engine(preprocess)
getattr(engine, entrypoint)(raw_request)
assert preprocess_calls == [raw_request]
assert engine.scheduler._waiting_queue == [prepared_request]
@pytest.mark.cpu
@pytest.mark.asyncio
async def test_async_add_req_and_stream_response():
@@ -123,47 +123,6 @@ def test_init_accepts_custom_scheduler(monkeypatch: pytest.MonkeyPatch) -> None:
assert engine.scheduler is custom_scheduler
def test_scheduler_initialization_failure_closes_scheduler_and_executor(monkeypatch: pytest.MonkeyPatch) -> None:
od_config = SimpleNamespace(
custom_pipeline_args=None,
model_class_name="SchedulerInitializationFailurePipeline",
streaming_output=False,
)
initialization_error = RuntimeError("Scheduler initialization failed")
custom_scheduler = SimpleNamespace(
initialize=Mock(side_effect=initialization_error),
close=Mock(),
)
fake_executor = SimpleNamespace(
shutdown=Mock(),
)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
lambda *args, **kwargs: None,
)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
lambda *args, **kwargs: None,
)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
lambda *args, **kwargs: Mock(return_value=fake_executor),
)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.supports_request_batch",
lambda *args, **kwargs: False,
)
with pytest.raises(RuntimeError) as exc_info:
DiffusionEngine(od_config, scheduler=custom_scheduler)
assert exc_info.value is initialization_error
custom_scheduler.initialize.assert_called_once_with(od_config)
custom_scheduler.close.assert_called_once_with()
fake_executor.shutdown.assert_called_once_with()
@pytest.mark.asyncio
async def test_step_compatibility_wrapper_returns_final_batch() -> None:
engine = _make_engine()
Loaded 100 of 237 files, more files were not shown because too many files have changed in this diff. Show more