Compare commits

...
Author SHA1 Message Date
Hongsheng Liu 7ec78360bb Merge branch 'main' into codex-split-attention-offload-guides 2026-08-15 07:00:22 +08:00
4b6d6d5cf7 [Bugfix][XPU][Tests]Scope XPU pytest to explicit paths to fix collection crash (#6175)
Signed-off-by: Joshna Medisetty <joshna.medisetty@intel.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Chendi.Xue <chendi.xue@intel.com>
2026-08-14 15:23:41 -05:00
Alex Brooks f7a5146093 [Feat/Bugfix] Fix API Server model tag <-> model sync & flag normalization (#3805)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-08-14 19:13:00 +00:00
Wenjia Li 8ecd1f6d5c [Perf][NPU] Fuse MiniMax H3 Qwen3-VL SwiGLU (#6167)
Signed-off-by: Wenjia Li <wjialish@gmail.com>
2026-08-14 11:47:17 +00:00
Wenjia Li ffd11c181e [Perf][NPU] Fuse MiniMax H3 Qwen3-VL RoPE (#6061)
Signed-off-by: Wenjia Li <wjialish@gmail.com>
2026-08-14 09:49:26 +00:00
R0CKSTAR 728640f4d9 [Bugfix][MUSA][Qwen Image] Avoid complex RoPE alias guards (#6110)
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
2026-08-14 15:51:23 +08:00
R0CKSTAR 022d7a7695 [Perf][Diffusion] Avoid redundant MiniMax-H3 reference video scans (#6064)
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
2026-08-14 15:49:26 +08:00
baonudesifeizhai b300f66656 [Bugfix] Fix HSDP compatibility with the new online FP8 linear method (#5677)
Signed-off-by: roG0d <baonudesifeizhai@gmail.com>
2026-08-14 15:45:55 +08:00
TJianandHongsheng Liu bc73546c06 [ROCm] [CI] Migrate CI to mi300x for v0.27.x (#5886)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
Signed-off-by: TJian <tunjian.tan@embeddedllm.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-14 15:38:47 +08:00
2b08419d8b [CI] Per-model, per-entry-mode code coverage collection (#5593)
Signed-off-by: Shenglei Fu <sfu@confluent.io>
Signed-off-by: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com>
Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>
Co-authored-by: wangyu <53896905+yenuo26@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-14 15:36:56 +08:00
zhengjiaandCursor 5d50bdc2ab [Perf] Cap diffusion worker thread count before spawn (#6165)
Signed-off-by: zjli2013 <leezhengjiang@126.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-14 12:14:55 +08:00
Samit 8a5bf2fb3f [doc]: Update latest news section for verl-omni release (#6187)
Signed-off-by: Samit <285365963@qq.com>
2026-08-14 12:04:15 +08:00
596c16a550 [Kernel] Fuse Q/K RMSNorm and RoPE (#5990)
Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Signed-off-by: Mengjie Zhao <zmj0129@gmail.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-08-14 00:46:59 +00:00
Lucas(Ruijie) 90a08c0eda [Model][TTS][WIP] Add dots.tts (rednote-hilab) — continuous-AR 48kHz TTS (#4765)
Signed-off-by: Moore-Z <zhumoore9661996@gmail.com>
2026-08-13 17:58:33 +00:00
Sy03 dbc0dd6d85 [Bugfix][Higgs-Audio-V3] Disable XQA decode (#6068)
Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: Sysy <1370724210@qq.com>
2026-08-13 10:00:50 -07:00
Mu GuanLin 83d1f11449 [Feat] LTX Standard Two-Stage Pipeline (#5500)
Signed-off-by: mglyn <1203789601@qq.com>
2026-08-14 00:13:14 +08:00
Hongsheng Liu 5569abf0a2 Update docs/design/feature/offloader/README.md
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 23:37:08 +08:00
Hongsheng Liu 10853fb4f3 Update docs/design/feature/offloader/README.md
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 23:37:01 +08:00
Yukim1andHongsheng Liu 4b66f62e10 [Refactor] Remove internal stage_configs_path plumbing (#5741)
Signed-off-by: zwhzzz0821 <2831474076@qq.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 21:29:03 +08:00
Weiming Liao f41451f64c [NPU] Upgrade to v0.27.0 (#6096)
Signed-off-by: Weiming Liao <liaowm5@gmail.com>
2026-08-13 20:01:01 +08:00
Hongsheng Liu bbe6ccc512 [Doc] Refresh community documentation (#6141)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 17:04:06 +08:00
Canlin Guo bbb709cfb5 [BugFix] Use MOSS-TTS-Local official sample params (#6156)
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
2026-08-13 08:56:26 +00:00
vraitiandClaude Opus 4.6 c98517ce7f [Metrics] Add Diffusion Metrics (#4755)
Signed-off-by: vraiti <vraiti@redhat.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-08-13 06:59:13 +00:00
NATUREandAlicia d77a35a32b [Bugfix] Scope Bagel FP8 config to diffusion stage (#6085)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-08-13 11:23:47 +08:00
Hongsheng Liu a52c19ea8b docs: organize offloader feature designs
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 11:35:07 +08:00
Hongsheng Liu 8b61499662 docs: rename offloader guides
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 09:30:34 +08:00
Hongsheng Liu 7872663bf0 docs: split diffusion feature guides
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 08:55:41 +08:00
186 changed files with 11196 additions and 2200 deletions

No files matched your search

+91 -141
View File
@@ -1,153 +1,103 @@
#!/bin/bash
# vllm-omni customized version
# Based on: vllm/.buildkite/scripts/hardware_ci/run-amd-test.sh
# Last synced: 2025-12-15
# Modifications: docker image name for vllm-omni
# This script runs test inside the corresponding ROCm docker container.
set -o pipefail
# Run vLLM-Omni ROCm tests directly in the MI300 Kubernetes pod. The pod's
# container image is selected by test-template-amd-omni.j2; MI300 has no DinD.
set -euo pipefail
# Export Python path
export PYTHONPATH=".."
: "${PYTHONFAULTHANDLER:=1}"
: "${HF_HOME:=/home/buildkite-agent/huggingface}"
: "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
: "${HF_HUB_ETAG_TIMEOUT:=60}"
: "${MIOPEN_DEBUG_CONV_DIRECT:=0}"
: "${MIOPEN_DEBUG_CONV_GEMM:=0}"
: "${VLLM_ROCM_USE_AITER:=0}"
export PYTHONFAULTHANDLER HF_HOME HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
export MIOPEN_DEBUG_CONV_DIRECT MIOPEN_DEBUG_CONV_GEMM VLLM_ROCM_USE_AITER
export PYTORCH_ROCM_ARCH=""
export PYTHONPATH="${PYTHONPATH:-..}"
if [[ "${VLLM_CI_DOCKER_DISABLED:-0}" != "1" ]]; then
echo "Error: MI300 CI must run natively with Docker disabled." >&2
exit 1
fi
if [[ -n "${TEST_COMMANDS:-}" ]]; then
commands="${TEST_COMMANDS}"
else
commands="$*"
fi
if [[ -z "${commands}" ]]; then
echo "Error: No test commands provided." >&2
exit 1
fi
echo "--- Native in-pod ROCm CI"
job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-local}}"
job_id="${job_id//[^A-Za-z0-9_.-]/_}"
native_cache_root="/tmp/vllm-omni-native-${job_id}"
TMPDIR="${native_cache_root}/tmp"
TORCHINDUCTOR_CACHE_DIR="${native_cache_root}/torchinductor"
TRITON_CACHE_DIR="${native_cache_root}/triton"
VLLM_CACHE_ROOT="${native_cache_root}/vllm"
XDG_CACHE_HOME="${native_cache_root}/xdg"
HF_DATASETS_CACHE="${native_cache_root}/huggingface/datasets"
export TMPDIR TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT
export XDG_CACHE_HOME HF_DATASETS_CACHE
mkdir -p "${TMPDIR}" "${TORCHINDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}" \
"${VLLM_CACHE_ROOT}" "${XDG_CACHE_HOME}" "${HF_HOME}" \
"${HF_DATASETS_CACHE}"
if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then
if ! command -v findmnt >/dev/null 2>&1; then
echo "Error: findmnt is required to verify the Hugging Face cache mount." >&2
exit 1
fi
hf_mount=$(findmnt -n -T "${HF_HOME}" -o TARGET 2>/dev/null || true)
if [[ -z "${hf_mount}" || "${hf_mount}" == "/" ]]; then
echo "Error: MI300 CI requires a persistent volume at or above ${HF_HOME}." >&2
exit 1
fi
fi
# Print ROCm version
echo "--- ROCm info"
rocminfo
# cleanup older docker images
cleanup_docker() {
# Get Docker's root directory
docker_root=$(docker info -f '{{.DockerRootDir}}')
if [ -z "$docker_root" ]; then
echo "Failed to determine Docker root directory."
exit 1
fi
echo "Docker root directory: $docker_root"
# Check disk usage of the filesystem where Docker's root directory is located
disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
# Define the threshold
threshold=70
if [ "$disk_usage" -gt "$threshold" ]; then
echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
# Remove dangling images (those that are not tagged and not used by any container)
docker image prune -f
# Remove unused volumes / force the system prune for old images as well.
docker volume prune -f && docker system prune --force --filter "until=72h" --all
echo "Docker images and volumes cleanup completed."
else
echo "Disk usage is below $threshold%. No cleanup needed."
fi
}
expected_gpus="${VLLM_CI_EXPECTED_GPU_COUNT:-1}"
python3 - <<'PY'
import os
# Call the cleanup docker function
cleanup_docker
import torch
echo "--- Pulling container"
## Temporary change to use AMD Docker Hub to store the vllm-omni image
# to bypass the rate limit issue with ECR Public Gallery.
# Images are now stored in a separate repository for vllm-omni, instead of vllm-ci.
# TODO: @tjtanaa point back to ECR Public Gallery
# once the amd agents are configured to use ECR Public Gallery.
# image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
image_name="rocm/vllm-omni:${BUILDKITE_COMMIT}"
container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
expected = int(os.environ.get("VLLM_CI_EXPECTED_GPU_COUNT", "1"))
assert torch.version.hip, "PyTorch is not a ROCm build"
assert torch.cuda.is_available(), "ROCm GPU is not available to PyTorch"
actual = torch.cuda.device_count()
assert actual == expected, f"Expected {expected} ROCm GPU(s), found {actual}"
PY
# TODO: @tjtanaa uncomment this once the amd agents are configured to use ECR Public Gallery.
# # Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
# sudo apt-get update && sudo apt-get install -y awscli
## Use safe docker login helper to prevent race conditions
# source "$(dirname "${BASH_SOURCE[0]}")/../../common/scripts/docker_login_ecr_public.sh"
# safe_docker_login_ecr_public
## Pull the container from AMD Docker Hub
docker pull "${image_name}"
remove_docker_container() {
docker ps -aq --filter "name=${container_name}" | xargs -r docker rm -f || true
docker image rm -f "${image_name}" || true
}
trap remove_docker_container EXIT
echo "--- Running container"
HF_CACHE="$(realpath ~)/huggingface"
mkdir -p "${HF_CACHE}"
HF_MOUNT="/root/.cache/huggingface"
if [[ -n "${TEST_COMMAND:-}" ]]; then
commands="$TEST_COMMAND"
else
commands="$@"
fi
echo "Commands:$commands"
PARALLEL_JOB_COUNT=8
MYPYTHONPATH=".."
# Test that we're launching on the machine that has
# proper access to GPUs
render_gid=$(getent group render | cut -d: -f3)
if [[ -z "$render_gid" ]]; then
echo "Error: 'render' group not found. This is required for GPU access." >&2
exit 1
fi
# check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs.
# TODO: @tjtanaa reenable to run VLLM_ROCM_USE_AITER=1 when AITER is shipped with prebuilt kernels.
if [[ $commands == *"--shard-id="* ]]; then
# assign job count as the number of shards used
commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g')
for GPU in $(seq 0 $(($PARALLEL_JOB_COUNT-1))); do
# assign shard-id for each shard
commands_gpu=$(echo "$commands" | sed -E "s/--shard-id[[:blank:]]*=[[:blank:]]*[0-9]*/--shard-id=${GPU} /g" | sed 's/ \\ / /g')
echo "Shard ${GPU} commands:$commands_gpu"
echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
docker run \
--device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
--network=host \
--shm-size=16gb \
--group-add "$render_gid" \
-e MIOPEN_DEBUG_CONV_DIRECT=0 \
-e MIOPEN_DEBUG_CONV_GEMM=0 \
-e VLLM_ROCM_USE_AITER=0 \
-e HIP_VISIBLE_DEVICES="${GPU}" \
-e HF_TOKEN \
-v "${HF_CACHE}:${HF_MOUNT}" \
-e "HF_HOME=${HF_MOUNT}" \
-e "PYTHONPATH=${MYPYTHONPATH}" \
--name "${container_name}_${GPU}" \
"${image_name}" \
/bin/bash -c "${commands_gpu}" \
|& while read -r line; do echo ">>Shard $GPU: $line"; done &
PIDS+=($!)
done
#wait for all processes to finish and collect exit codes
for pid in "${PIDS[@]}"; do
wait "${pid}"
STATUS+=($?)
done
for st in "${STATUS[@]}"; do
if [[ ${st} -ne 0 ]]; then
echo "One of the processes failed with $st"
exit "${st}"
echo "Commands:${commands}"
if /bin/bash -o pipefail -c '
set -E
test_status=0
trap '\''
command_status=$?
# Keep running subsequent commands, but retain a failure for the final exit.
# Prefer a test failure over pytest "no tests collected" when both occur.
if (( test_status == 0 || (test_status == 5 && command_status != 5) )); then
test_status=${command_status}
fi
done
'\'' ERR
eval "$1"
exit "${test_status}"
' _ "${commands}"; then
exit 0
else
echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
docker run \
--device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
--network=host \
--shm-size=16gb \
--group-add "$render_gid" \
-e MIOPEN_DEBUG_CONV_DIRECT=0 \
-e MIOPEN_DEBUG_CONV_GEMM=0 \
-e VLLM_ROCM_USE_AITER=0 \
-e HF_TOKEN \
-v "${HF_CACHE}:${HF_MOUNT}" \
-e "HF_HOME=${HF_MOUNT}" \
-e "PYTHONPATH=${MYPYTHONPATH}" \
--name "${container_name}" \
"${image_name}" \
/bin/bash -c "${commands}"
exit_code=$?
fi
if [[ ${exit_code} -eq 5 && "${VLLM_CI_ALLOW_NO_TESTS:-0}" == "1" ]]; then
echo "Pytest collected no tests; VLLM_CI_ALLOW_NO_TESTS=1, treating exit code 5 as success."
exit 0
fi
exit "${exit_code}"
+45 -34
View File
@@ -7,54 +7,64 @@ steps:
- group: ":card_index_dividers: Simple Test"
steps:
- label: "Simple · Diffusion & Model Executor Test (1/2)"
agent_pool: mi325_1
- label: "Simple · Model Executor Test · Shard %N/%t"
agent_pool: mi300_1
parallelism: 2
depends_on: amd-build
timeout_in_minutes: 70
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
# CPU tests on ROCm gpu machines are very slow, so we give them more time.
- "timeout 60m pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- "pytest -sv tests/model_executor -m 'core_model and cpu' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:coverage-model-executor-$$BUILDKITE_PARALLEL_JOB.xml"
- label: "Simple · Diffusion & Model Executor Test (2/2)"
agent_pool: mi325_1
- label: "Simple · Diffusion Test · Shard %N/%t"
agent_pool: mi300_1
parallelism: 2
depends_on: amd-build
timeout_in_minutes: 70
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# CPU tests on ROCm gpu machines are very slow, so we give them more time.
- "timeout 60m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:coverage-diffusion-$$BUILDKITE_PARALLEL_JOB.xml"
- label: "Simple · Engine&Entrypoints Test"
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
- "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Simple · Other Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
timeout_in_minutes: 50
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- group: ":card_index_dividers: Diffusion Test"
steps:
- label: "Diffusion · LoRA Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
artifact_paths:
- lora-server.log
commands:
- timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
- pytest -sv tests/diffusion/lora/ -m 'advanced_model and cuda' --run-level "advanced_model" 2>&1 | tee lora-server.log
- label: "Diffusion Model CPU offloading Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -62,7 +72,7 @@ steps:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
timeout 20m bash -c '
bash -c '
set +e
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
EXIT1=\$?
@@ -72,7 +82,7 @@ steps:
'
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
agent_pool: mi325_4
agent_pool: mi300_4
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -84,7 +94,7 @@ steps:
# merge-only tests
- label: "Diffusion Tensor Parallelism Test"
agent_pool: mi325_2
agent_pool: mi300_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -92,7 +102,7 @@ steps:
- pytest -s -v tests/diffusion/distributed/test_tensor_parallel.py
- label: "Diffusion Model Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -100,7 +110,7 @@ steps:
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py
- label: "Diffusion GPU Worker Test"
agent_pool: mi325_2
agent_pool: mi300_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -108,7 +118,7 @@ steps:
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
- label: "Diffusion · Wan22 Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -121,7 +131,7 @@ steps:
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
# - label: "Audio Generation Model Test"
# agent_pool: mi325_1
# agent_pool: mi300_1
# depends_on: amd-build
# mirror_hardwares: [amdproduction]
# grade: Blocking
@@ -132,7 +142,7 @@ steps:
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
- label: "Engine Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -147,7 +157,7 @@ steps:
- group: ":card_index_dividers: E2E Test"
steps:
- label: "Omni Model Test Qwen2-5-Omni"
agent_pool: mi325_4
agent_pool: mi300_4
depends_on: amd-build
mirror_hardwares: [amdproduction]
commands:
@@ -155,7 +165,7 @@ steps:
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
- label: "Omni Model Test Qwen3-Omni"
agent_pool: mi325_2
agent_pool: mi300_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -163,22 +173,23 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
- label: "Qwen3-TTS CustomVoice E2E Test"
agent_pool: mi325_1
- label: "Qwen3-TTS CustomVoice E2E Test · Shard %N/%t"
agent_pool: mi300_1
parallelism: 2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- |
timeout 50m bash -c '
bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
'
- label: "Qwen3-TTS Base E2E Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -196,7 +207,7 @@ steps:
# Need to debug before reneable numerical changes across large PRs
# # split Bagel Model Test with H100 (Real Weights) into three tests
# - label: "Bagel Text2Img Model Test (1/3)"
# agent_pool: mi325_1
# agent_pool: mi300_1
# depends_on: amd-build
# mirror_hardwares: [amdproduction]
# grade: Blocking
@@ -208,7 +219,7 @@ steps:
# - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
# - label: "Bagel Img2Img Model Test (2/3)"
# agent_pool: mi325_1
# agent_pool: mi300_1
# depends_on: amd-build
# mirror_hardwares: [amdproduction]
# grade: Blocking
@@ -220,7 +231,7 @@ steps:
# - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
# - label: "Bagel Online Serving Test (3/3)"
# agent_pool: mi325_1
# agent_pool: mi300_1
# depends_on: amd-build
# mirror_hardwares: [amdproduction]
# grade: Blocking
+47 -34
View File
@@ -7,8 +7,9 @@ steps:
- group: ":card_index_dividers: Simple Test"
steps:
- label: "Simple · Diffusion & Model Executor Test"
agent_pool: mi325_1
- label: "Simple · Model Executor Test · Shard %N/%t"
agent_pool: mi300_1
parallelism: 2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -16,10 +17,30 @@ steps:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- "pytest -sv tests/model_executor -m 'core_model and cpu' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:coverage-model-executor-$$BUILDKITE_PARALLEL_JOB.xml"
- label: "Simple · Diffusion Test"
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Simple · Engine&Entrypoints Test"
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_ROCM_USE_AITER=0
- "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Simple · Other Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -27,19 +48,19 @@ steps:
- export VLLM_ROCM_USE_AITER=0
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
# that is not supported on CPU
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
- label: "Custom Pipeline Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model"
- pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model"
- label: "Voxtral TTS CUDA Unit Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -49,33 +70,25 @@ steps:
- group: ":card_index_dividers: Diffusion Test"
steps:
- label: "Diffusion · Batch Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
- pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
- label: "Diffusion · Offloader Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
EXIT1=\$?
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
EXIT2=\$?
exit \$((EXIT1 | EXIT2))
'
- pytest -s -v tests/diffusion/offloader -m "core_model and cuda" --run-level "core_model"
- label: "Diffusion · Cache Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -83,10 +96,10 @@ steps:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
- pytest -sv tests/diffusion/cache -m 'core_model and cuda' --run-level "core_model"
- label: "Diffusion Sequence Parallelism Test"
agent_pool: mi325_2
agent_pool: mi300_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -96,7 +109,7 @@ steps:
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
- label: "Diffusion GPU Worker Test"
agent_pool: mi325_2
agent_pool: mi300_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -104,7 +117,7 @@ steps:
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
- label: "Diffusion · Model Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -112,7 +125,7 @@ steps:
- timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
- label: "Diffusion · Wan22 Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -126,7 +139,7 @@ steps:
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
# - label: "Audio Generation Model Test"
# agent_pool: mi325_1
# agent_pool: mi300_1
# depends_on: amd-build
# mirror_hardwares: [amdproduction]
# grade: Blocking
@@ -137,7 +150,7 @@ steps:
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
- label: "Engine Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -149,7 +162,7 @@ steps:
'
# - label: "Omni Model Test Qwen3-Omni"
# agent_pool: mi325_2
# agent_pool: mi300_2
# depends_on: amd-build
# mirror_hardwares: [amdproduction]
# grade: Blocking
@@ -162,7 +175,7 @@ steps:
- group: ":card_index_dividers: E2E Test"
steps:
- label: "Qwen3-TTS E2E Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -173,7 +186,7 @@ steps:
- timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
- label: "CosyVoice3-TTS E2E Test"
agent_pool: mi325_1
agent_pool: mi300_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
@@ -186,7 +199,7 @@ steps:
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
# Need to debug before reneable numerical changes across large PRs
# - label: "Bagel Text2Img Model Test"
# agent_pool: mi325_1
# agent_pool: mi300_1
# depends_on: amd-build
# mirror_hardwares: [amdproduction]
# grade: Blocking
@@ -198,7 +211,7 @@ steps:
# - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"
# - label: "Bagel Img2Img Model Test"
# agent_pool: mi325_1
# agent_pool: mi300_1
# depends_on: amd-build
# mirror_hardwares: [amdproduction]
# grade: Blocking
@@ -210,7 +223,7 @@ steps:
# - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"
# - label: "Bagel Online Serving Test"
# agent_pool: mi325_1
# agent_pool: mi300_1
# depends_on: amd-build
# mirror_hardwares: [amdproduction]
# grade: Blocking
+104 -8
View File
@@ -1,11 +1,65 @@
{# vllm-omni customized version
Based on: https://github.com/vllm-project/ci-infra/blob/main/buildkite/test-template-amd.j2
Last synced: 2025-12-15
Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests
Last synced: 2026-08-04
Modifications: Use the shared MI300 queues and Omni command-passing contract
#}
{% set docker_image_amd = "rocm/vllm-omni:$BUILDKITE_COMMIT" %}
{% set default_working_dir = "/app/vllm-omni" %}
{# MI300 runs tests natively in Kubernetes; Docker-in-Docker is unavailable. #}
{% macro amd_native_k8s_plugin(indent, step) -%}
{{ indent }}plugins:
{{ indent }} - kubernetes:
{{ indent }} podSpecPatch:
{{ indent }} automountServiceAccountToken: false
{{ indent }} securityContext:
{{ indent }} seccompProfile:
{{ indent }} type: RuntimeDefault
{{ indent }} imagePullSecrets:
{{ indent }} - name: docker-config
{{ indent }} containers:
{{ indent }} - name: container-0
{{ indent }} image: {{ docker_image_amd }}
{{ indent }} imagePullPolicy: Always
{{ indent }} securityContext:
{{ indent }} allowPrivilegeEscalation: false
{{ indent }} capabilities:
{{ indent }} add:
{{ indent }} - IPC_LOCK
{{ indent }} resources:
{{ indent }} limits:
{{ indent }} amd.com/gpu: "{{ step.agent_pool[-1:] }}"
{{ indent }} requests:
{{ indent }} amd.com/gpu: "{{ step.agent_pool[-1:] }}"
{{ indent }} volumeMounts:
{{ indent }} - name: devshm
{{ indent }} mountPath: /dev/shm
{{ indent }} env:
{{ indent }} - name: AMD_CI_RUNTIME
{{ indent }} value: native
{{ indent }} - name: NATIVE_CI
{{ indent }} value: "true"
{{ indent }} - name: VLLM_CI_DOCKER_DISABLED
{{ indent }} value: "1"
{{ indent }} - name: VLLM_CI_EXPECTED_GPU_COUNT
{{ indent }} value: "{{ step.agent_pool[-1:] }}"
{{ indent }} - name: HF_HOME
{{ indent }} value: /home/buildkite-agent/huggingface
{{ indent }} - name: MIOPEN_DEBUG_CONV_DIRECT
{{ indent }} value: "0"
{{ indent }} - name: MIOPEN_DEBUG_CONV_GEMM
{{ indent }} value: "0"
{{ indent }} - name: VLLM_ROCM_USE_AITER
{{ indent }} value: "0"
{{ indent }} - name: PYTORCH_ROCM_ARCH
{{ indent }} value: ""
{{ indent }} volumes:
{{ indent }} - name: devshm
{{ indent }} emptyDir:
{{ indent }} medium: Memory
{{ indent }} sizeLimit: 16Gi
{%- endmacro %}
{% if env %}
env:
{% for name in env %}
@@ -53,19 +107,31 @@ steps:
{% if step.agent_pool %}
queue: amd_{{ step.agent_pool }}
{% else %}
queue: amd_mi325_1
queue: amd_mi300_1
{% endif %}
{{ amd_native_k8s_plugin(" ", step) }}
{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
{% set indented_cmd = cmd_body | replace("\n", "\n ") %}
command: bash .buildkite/amd/scripts/run-amd-test.sh
env:
DOCKER_BUILDKIT: "1"
AMD_CI_RUNTIME: "native"
NATIVE_CI: "true"
VLLM_CI_DOCKER_DISABLED: "1"
VLLM_CI_EXPECTED_GPU_COUNT: "{{ step.agent_pool[-1:] }}"
VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE: "1"
HF_HOME: "/home/buildkite-agent/huggingface"
HF_HUB_DOWNLOAD_TIMEOUT: "300"
HF_HUB_ETAG_TIMEOUT: "60"
MIOPEN_DEBUG_CONV_DIRECT: "0"
MIOPEN_DEBUG_CONV_GEMM: "0"
VLLM_ROCM_USE_AITER: "0"
PYTORCH_ROCM_ARCH: ""
{% if step.env %}
{% for name in step.env %}
{{ name }}: "{{ step.env[name] }}"
{% endfor %}
{% endif %}
TEST_COMMAND: |-
TEST_COMMANDS: |-
(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
{% if env %}
{% for name in env %}
@@ -82,6 +148,15 @@ steps:
{% endif %}
{{ indented_cmd | safe }}
priority: 100
{% if step.parallelism %}
parallelism: {{ step.parallelism }}
{% endif %}
{% if step.artifact_paths %}
artifact_paths:
{% for path in step.artifact_paths %}
- "{{ path }}"
{% endfor %}
{% endif %}
{% if step.timeout_in_minutes %}
timeout_in_minutes: {{ step.timeout_in_minutes }}
{% endif %}
@@ -100,19 +175,31 @@ steps:
{% if step.agent_pool %}
queue: amd_{{ step.agent_pool }}
{% else %}
queue: amd_mi325_1
queue: amd_mi300_1
{% endif %}
{{ amd_native_k8s_plugin(" ", step) }}
{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
{% set indented_cmd = cmd_body | replace("\n", "\n ") %}
command: bash .buildkite/amd/scripts/run-amd-test.sh
env:
DOCKER_BUILDKIT: "1"
AMD_CI_RUNTIME: "native"
NATIVE_CI: "true"
VLLM_CI_DOCKER_DISABLED: "1"
VLLM_CI_EXPECTED_GPU_COUNT: "{{ step.agent_pool[-1:] }}"
VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE: "1"
HF_HOME: "/home/buildkite-agent/huggingface"
HF_HUB_DOWNLOAD_TIMEOUT: "300"
HF_HUB_ETAG_TIMEOUT: "60"
MIOPEN_DEBUG_CONV_DIRECT: "0"
MIOPEN_DEBUG_CONV_GEMM: "0"
VLLM_ROCM_USE_AITER: "0"
PYTORCH_ROCM_ARCH: ""
{% if step.env %}
{% for name in step.env %}
{{ name }}: "{{ step.env[name] }}"
{% endfor %}
{% endif %}
TEST_COMMAND: |-
TEST_COMMANDS: |-
(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
{% if env %}
{% for name in env %}
@@ -129,6 +216,15 @@ steps:
{% endif %}
{{ indented_cmd | safe }}
priority: 100
{% if step.parallelism %}
parallelism: {{ step.parallelism }}
{% endif %}
{% if step.artifact_paths %}
artifact_paths:
{% for path in step.artifact_paths %}
- "{{ path }}"
{% endfor %}
{% endif %}
{% if step.timeout_in_minutes %}
timeout_in_minutes: {{ step.timeout_in_minutes }}
{% endif %}
@@ -9,6 +9,11 @@
# 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
#
# The ``l4_*`` presets run under the ``docker`` plugin, which does not put
# ``buildkite-agent`` on PATH inside the container the way the ``kubernetes``
# presets do, so they set ``mount-buildkite-agent: true``; without it any step
# calling ``buildkite-agent`` (artifact upload, annotate, meta-data) fails there.
x-npu-k8s-plugins: &npu_k8s_plugins
- kubernetes:
@@ -26,6 +31,7 @@ mirror_hardwares:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
mount-buildkite-agent: true
shm-size: "8gb"
environment:
- HF_HOME=/fsx/hf_cache
@@ -41,6 +47,7 @@ mirror_hardwares:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
mount-buildkite-agent: true
shm-size: "8gb"
environment:
- HF_HOME=/fsx/hf_cache
+82
View File
@@ -0,0 +1,82 @@
#!/usr/bin/env bash
# Run a model's E2E tests once per entry mode as separate coverage runs and upload
# the reports, so per-mode attribution survives (see vllm-project/vllm-omni#5332).
#
# Usage: run_cov_split.sh --model-id <id> [--offline <paths>] [--online <paths>] \
# -- <pytest args...>
#
# --model-id Names the artifacts: coverage-<id>-{offline,online}-<step id>.xml.gz.
# Use the model's directory name under vllm_omni/*/models/.
# --offline Pytest path(s) for the offline half. Repeatable.
# --online Pytest path(s) for the online half. Repeatable.
# -- Everything after this is forwarded verbatim to every pytest run,
# quoting preserved: the job's own -m expression, --run-level, and
# anything else it needs (e.g. --test-config-file).
#
# At least one of --offline/--online is required, so a job with only one mode runs
# just that half. Total runtime is bounded by the step's timeout_in_minutes.
#
# Exits non-zero if any half or the upload failed.
set -uo pipefail
MODEL_ID=""
OFFLINE=()
ONLINE=()
PYTEST_ARGS=()
while [[ $# -gt 0 ]]; do
case "$1" in
--model-id) MODEL_ID="$2"; shift 2 ;;
--offline) OFFLINE+=("$2"); shift 2 ;;
--online) ONLINE+=("$2"); shift 2 ;;
--) shift; PYTEST_ARGS=("$@"); break ;;
*) echo "run_cov_split.sh: unknown argument: $1" >&2; exit 2 ;;
esac
done
if [[ -z "${MODEL_ID}" ]]; then
echo "run_cov_split.sh: missing --model-id" >&2
exit 2
fi
if ((${#OFFLINE[@]} == 0 && ${#ONLINE[@]} == 0)); then
echo "run_cov_split.sh: need at least one of --offline / --online" >&2
exit 2
fi
# Two steps can test the same model, and artifacts are scoped to the build (a
# nightly one carries the ready, merge and nightly tiers at once), so unsuffixed
# reports would share a path -- a glob download then races to rename onto one file.
# Not BUILDKITE_JOB_ID: the step id survives retries, so attempts collapse under
# Buildkite's latest-attempt lookup instead of accumulating. A parallel step would
# need BUILDKITE_PARALLEL_JOB as well; no pilot is one.
REPORT_SUFFIX="${BUILDKITE_STEP_ID:-local}"
OFFLINE_REPORT="coverage-${MODEL_ID}-offline-${REPORT_SUFFIX}.xml"
ONLINE_REPORT="coverage-${MODEL_ID}-online-${REPORT_SUFFIX}.xml"
run_half() {
local mode="$1"
shift
local report="coverage-${MODEL_ID}-${mode}-${REPORT_SUFFIX}.xml"
echo "--- coverage: ${MODEL_ID} ${mode}"
pytest -s -v "$@" "${PYTEST_ARGS[@]}" --cov=vllm_omni --cov-report="xml:${report}"
local status=$?
# Cobertura lists every statement in vllm_omni whether or not it ran, so the
# report is ~7MB no matter how little the job covers. The repetition is what
# dominates it, and gzip takes that to ~400KB.
gzip -9 -f "${report}" || status=1
return "${status}"
}
# Clear both modes, not just the ones being run: the upload glob matches both, so
# a report left by an earlier run would otherwise be published as if the mode it
# belongs to had run this time.
rm -f "${OFFLINE_REPORT}" "${ONLINE_REPORT}" "${OFFLINE_REPORT}.gz" "${ONLINE_REPORT}.gz"
EXIT=0
((${#OFFLINE[@]})) && { run_half offline "${OFFLINE[@]}" || EXIT=1; }
((${#ONLINE[@]})) && { run_half online "${ONLINE[@]}" || EXIT=1; }
buildkite-agent artifact upload "coverage-${MODEL_ID}-*-${REPORT_SUFFIX}.xml.gz" || EXIT=1
exit "${EXIT}"
+21 -7
View File
@@ -94,13 +94,19 @@ steps:
- vllm_omni/model_executor/models/qwen3_tts/
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
- vllm_omni/deploy/qwen3_tts.yaml
- .buildkite/common/scripts/run_cov_split.sh
- pyproject.toml
# 40 -> 60: the coverage split runs the model twice.
timeout_in_minutes: 60
commands:
- |
timeout 40m bash -c "
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model'
"
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
.buildkite/common/scripts/run_cov_split.sh \
--model-id qwen3_tts \
--offline tests/e2e/offline_inference/test_qwen3_tts_base.py \
--online tests/e2e/online_serving/test_qwen3_tts_base.py \
-- -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: l4_1
- label: "TTS · Higgs-Audio-V3 Test"
@@ -203,9 +209,17 @@ steps:
- vllm_omni/diffusion/models/bagel/
- vllm_omni/model_executor/stage_input_processors/bagel.py
- vllm_omni/deploy/bagel.yaml
timeout_in_minutes: 20
- .buildkite/common/scripts/run_cov_split.sh
- pyproject.toml
# 20 -> 40: the coverage split runs the model twice.
timeout_in_minutes: 40
commands:
- pytest -s -v tests/e2e/offline_inference/test_bagel.py tests/e2e/online_serving/test_bagel.py -m 'advanced_model and cuda' --run-level 'advanced_model'
- |
.buildkite/common/scripts/run_cov_split.sh \
--model-id bagel \
--offline tests/e2e/offline_inference/test_bagel.py \
--online tests/e2e/online_serving/test_bagel.py \
-- -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: h100_1
- label: "Diffusion · Qwen Image Test"
+8 -2
View File
@@ -456,7 +456,7 @@ steps:
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
- pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx2_expansion.py -m "full_model and cuda" --run-level "full_model"
- pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx.py -m "full_model and cuda" --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test"
@@ -477,10 +477,16 @@ steps:
commands:
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_one_stage_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model
mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2V · LTX Accuracy Test"
timeout_in_minutes: 180
commands:
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx23_two_stage_layer_fused_i2v]' -m full_model --run-level full_model
mirror_hardwares: h100_1
- label: ":full_moon: Diffusion X2V · Perf Test"
key: nightly-diffusion-x2v-performance
timeout_in_minutes: 180
+4 -3
View File
@@ -87,8 +87,9 @@ time timeout -k 30 30m docker run \
echo $ZE_AFFINITY_MASK
pip install tblib==3.1.0
cd /workspace/vllm-omni
pytest -v -s -m "core_model and xpu and B60"
XPU_TEST_PATHS="tests/diffusion tests/dfx tests/e2e"
pytest -v -s $XPU_TEST_PATHS -m "core_model and xpu and B60"
pytest -v -s tests/diffusion/quantization/test_mxfp8_config.py
pytest -v -s -m "advanced_model and xpu and B60"
pytest -v -s -m "omni and xpu and B60"
pytest -v -s $XPU_TEST_PATHS -m "advanced_model and xpu and B60"
pytest -v -s $XPU_TEST_PATHS -m "omni and xpu and B60"
'
+2
View File
@@ -46,6 +46,8 @@ htmlcov/
.cache
nosetests.xml
coverage.xml
coverage-*.xml
coverage-*.xml.gz
*.cover
*.py,cover
.hypothesis/
+1
View File
@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
---
*Latest News* 🔥
- [2026/08] [VeRL-Omni](https://github.com/verl-project/verl-omni) `v0.2.0` is released: faster diffusion RL powered by vLLM-Omni (request-level/step-wise batching with FA3), rebuilt Qwen3-Omni multimodal training (DPO & GSPO), plus LTX-2.3, Qwen-Image-Edit support and more. See the [release notes](https://github.com/verl-project/verl-omni/releases/tag/v0.2.0).
- [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support.
- [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
- [2026/06] Starting with [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0), vLLM-Omni publishes a stable release aligned with every even-numbered upstream vLLM minor version. [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0), [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0), [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0), and [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) continued this cadence, expanding omni and world-model support with [NVIDIA Cosmos3](recipes/cosmos3/Cosmos3-Nano.md) and DreamZero, adding models such as MiniCPM-o 4.5, MOSS-TTS, and Lance, and advancing TTS, diffusion, distributed execution, quantization, RL integration through [VeRL-Omni](https://github.com/verl-project/verl-omni), and CUDA/ROCm/MUSA/NPU/XPU coverage.
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.26.0
ARG VLLM_ASCEND_TAG=v0.27.0
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
# WORKDIR /vllm-workspace/vllm
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.26.0-a3
ARG VLLM_ASCEND_TAG=v0.27.0-a3
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
# WORKDIR /vllm-workspace/vllm
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.26.0
ARG VLLM_ASCEND_TAG=v0.27.0
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
ARG APP_DIR=/vllm-workspace/vllm-omni
+1 -1
View File
@@ -1,5 +1,5 @@
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.26.0-a3
ARG VLLM_ASCEND_TAG=v0.27.0-a3
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
ARG APP_DIR=/vllm-workspace/vllm-omni
+9 -2
View File
@@ -1,4 +1,4 @@
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.26.0
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.27.0
FROM ${BASE_IMAGE} AS base
# Declare a variable to know if we want to use the nightly build or the stable build.
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
# we should swap over to use stable release ASAP.
# We should avoid relying on custom commits.
ARG USE_NIGHTLY_BUILD=0
ARG VLLM_VERSION_OR_COMMIT_HASH=v0.26.0
ARG VLLM_VERSION_OR_COMMIT_HASH=v0.27.0
ARG ARG_PYTORCH_ROCM_ARCH
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
@@ -79,9 +79,16 @@ RUN printf '%s\n' \
' os._exit(_exit_code)' \
> ${COMMON_WORKDIR}/vllm-omni/conftest.py
# Some optional dependencies pull in the CPU ONNX Runtime distribution. The
# CPU and ROCm wheels own the same import path and cannot coexist, so remove
# the CPU wheel and reinstall the ROCm wheel after dependency resolution.
RUN cd ${COMMON_WORKDIR}/vllm-omni && \
uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" \
--no-cache-dir ".[dev]" --no-build-isolation && \
python3 -m pip uninstall -y onnxruntime && \
python3 -m pip install --no-cache-dir --force-reinstall --no-deps \
"onnxruntime-rocm>=1.22.2" && \
python3 -c "import importlib.metadata as m; import onnxruntime as ort; assert 'ROCMExecutionProvider' in ort.get_available_providers(), ort.get_available_providers(); assert m.version('onnxruntime-rocm'); assert not any(d.metadata['Name'] == 'onnxruntime' for d in m.distributions())" && \
python3 -c "from torchcodec.decoders import VideoDecoder"
RUN ln -sf /usr/bin/python3 /usr/bin/python
+20 -4
View File
@@ -60,7 +60,11 @@ nav:
- Overview and Compatibility:
- Overview: user_guide/diffusion_features.md
- Feature Compatibility: user_guide/feature_compatibility.md
- CPU Offloading: user_guide/diffusion/cpu_offload.md
- CPU Offloading:
- Overview: user_guide/diffusion/cpu_offload.md
- Model-Level: user_guide/diffusion/offloader/cpu_offload.md
- Layerwise: user_guide/diffusion/offloader/layerwise_offload.md
- Distributed Layerwise: user_guide/diffusion/offloader/distributed_layerwise_offload.md
- Cache Acceleration:
- TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
- Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
@@ -72,7 +76,13 @@ 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
- Attention Backends:
- Overview: user_guide/diffusion/attention_backends.md
- Dense Backends: user_guide/diffusion/attention_backends/dense_backends.md
- TRTLLM: user_guide/diffusion/attention_backends/trtllm.md
- SageAttention: user_guide/diffusion/attention_backends/sage.md
- Hugging Face Hub: user_guide/diffusion/attention_backends/huggingface_hub.md
- RainFusion: user_guide/diffusion/attention_backends/rainfusion.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
@@ -125,11 +135,16 @@ nav:
- Tensor Parallel: design/feature/tensor_parallel.md
- VAE Patch Parallelism: design/feature/vae_parallel.md
- Attention Optimization:
- Backend Selection: design/feature/attention_backend_selection.md
- Skip-Softmax: design/feature/skip_softmax.md
- CPU Offloading:
- Overview: design/feature/offloader/README.md
- Model-Level: design/feature/offloader/cpu_offload.md
- Layerwise: design/feature/offloader/layerwise_offload.md
- Distributed Layerwise: design/feature/offloader/distributed_layerwise_offload.md
- Cache-DiT: design/feature/cache_dit.md
- TeaCache: design/feature/teacache.md
- Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md
- Distributed Layerwise Offload: design/feature/distributed_layerwise_offload.md
- Infrastructure and Performance:
- Prometheus Metrics: design/metrics.md
- Speech Generation Performance Optimizations: design/qwen3_omni_tts_performance_optimization.md
@@ -163,7 +178,8 @@ nav:
- CLI Reference: cli
- Community:
- Governance: community/governance.md
- community/*
- Contact Us: community/contact_us.md
- Community Meetings: community/meetings.md
- Slack: https://slack.vllm.ai
- Blog: https://blog.vllm.ai
- Forum: https://discuss.vllm.ai
+7 -3
View File
@@ -1,5 +1,9 @@
# Contact Us
- For technical questions and feature requests, please use GitHub [Issues](https://github.com/vllm-project/vllm-omni/issues)
- For coordinating contributions and development and discussing with other users and developers, please join `sig-omni` channel in our [Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/)
- For security disclosures, please use GitHub's [Security Advisories](https://github.com/vllm-project/vllm-omni/security/advisories) feature
- For bugs and feature requests, search or open a GitHub [issue](https://github.com/vllm-project/vllm-omni/issues).
- For technical and usage questions, join the `sig-omni` channel in [vLLM Slack](https://slack.vllm.ai/) or use the [vLLM Forum](https://discuss.vllm.ai/).
- For coordinating contributions and development, use `sig-omni`, the forum, or join a [community meeting](meetings.md).
## Security Disclosures
Do not share vulnerability details in public GitHub issues, pull requests, discussions, or Slack channels. To arrange a private report, send a direct message to one of the [lead maintainers](governance.md#lead-maintainers) in [vLLM Slack](https://slack.vllm.ai/). If you cannot send a direct message, you may post a detail-free request in `sig-omni` asking a lead maintainer to contact you privately; do not include any vulnerability details in that request.
+3 -23
View File
@@ -52,31 +52,11 @@ Committers have write access and merge rights. They typically have deep expertis
- [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
- [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
## Diffusion Workload Division
## Reviewer Routing
This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review.
For path-based reviewer routing, consult [`.github/CODEOWNERS`](https://github.com/vllm-project/vllm-omni/blob/main/.github/CODEOWNERS). Detailed subsystem boundaries and safe-change guidance live in the [module design documents](../design/index.md#module-design-documents).
### Runtime-related
- Diffusion models: @RuixiangMa, @wtomin, @xuechendi
### Optimization
- Parallel: @RuixiangMa, @wtomin
- Attention backend: @ZJY0516
- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
- Offload: @yuanheng-zhao
- Intel GPU support: @xuechendi
- Quantization: @lishunyang12, @david6666666
- Kernel fusion / communication-computation: @ZJY0516
### Docs & Test
- Diffusion docs: @RuixiangMa, @wtomin
## Meetings
Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
`CODEOWNERS` is operational routing metadata, not a governance or approval policy. Being listed does not grant committer status or merge rights, and it does not mean every listed reviewer must approve a change. The maintainer roster above is authoritative for project roles and merge rights.
## Committer Nomination Process
+8
View File
@@ -0,0 +1,8 @@
# Community Meetings
vLLM-Omni hosts developer-facing meetings for Chinese- and English-language audiences. Schedules and access details can change, so use these maintained documents as the source of truth:
- [Chinese-language meeting document](https://docs.google.com/document/d/1pdUBiS_7mdOUNDtdwy-9OUf7jsMWN324BIbps5olDME/edit?tab=t.0)
- [English-language meeting document](https://docs.google.com/document/d/19AlD08T8uJ7wqfs8VtYCi0iGMTwlTYVVlAxFZ3S7OIg/edit?tab=t.0)
Each document contains the current schedule, access details, agendas, and past meeting notes. For broader vLLM community events, see the [vLLM events calendar](https://vllm.ai/events/).
-1
View File
@@ -1 +0,0 @@
# Meetups
-12
View File
@@ -1,12 +0,0 @@
# Volunteers for Bugfix and CI
We encourage you to check current docs and [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions for your questions. If non of these can solve it, please propose an issue to describe your questions about bug or CI problems for developing.
If you have urgent need for locating and solving bugfix or CI problems, please find community volunteers below.
| Dec 4-Dec 12 | Dec 15-Dec 19 | Dec 22-Dec 26 | Dec 29- Jan 2, 2026| Jan 5-Jan 9 | Jan 12-Jan 16 |
|----------|----------|----------|----------|----------|----------|
| <a href="https://github.com/congw729">Conw729</a> | <a href="https://github.com/yinpeiqi">yinpeiqi</a> | <a href="https://github.com/tzhouam">tzhouam</a> | <a href="https://github.com/SamitHuang">SamitHuang</a> | <a href="https://github.com/gcanlin">gcanlin</a> | <a href="https://github.com/natureofnature">natureofnature</a> |
| <a href="https://github.com/david6666666">david6666666</a> | <a href="https://github.com/R2-Y">R2-Y</a> | <a href="https://github.com/hsliuustc0106">hsliuustc0106</a> | <a href="https://github.com/Gaohan123">Gaohan123</a> | <a href="https://github.com/ZJY0516">ZJY0516</a> | <a href="https://github.com/qibaoyuan">qibaoyuan</a> |
We kindly welcome more contributors to fix bugs and contribute new features!
+2 -2
View File
@@ -3,7 +3,7 @@
Thank you for your interest in contributing to vLLM-Omni! This document provides guidelines and instructions for contributing.
!!! note
We host weekly developer-facing online meetings to discuss milestones and updates **every Tuesday at 19:30 PDT**. Meeting link as well as the past meeting notes can be found [here](https://tinyurl.com/vllm-omni-meeting).
vLLM-Omni hosts developer-facing meetings for Chinese- and English-language audiences. See [Community Meetings](../community/meetings.md) for current schedules, access details, agendas, and past notes.
## Getting Started
@@ -122,7 +122,7 @@ Repository skills guide the agent, but they do not replace contributor judgment,
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.
!!! important
If you discover a security vulnerability, please report it by creating a GitHub issue with the `security` label.
Do not report suspected security vulnerabilities through a public issue, pull request, discussion, or Slack channel. Follow the [security disclosure instructions](../community/contact_us.md#security-disclosures) to arrange a private report.
## Pull Requests & Code Reviews
+90
View File
@@ -327,6 +327,96 @@ On PR builds, `upload_pipeline.py` logs `skip '…' (no changes under …)` for
- Use **groups** for dashboard readability; keep **E2E Test** as the group name CUDA diff filtering expects for `--e2e` nightly runs on main.
- Prefix labels by model domain: **Omni ·**, **TTS ·**, **Diffusion ·**, **Simple ·**, etc., matching existing steps.
### Per-model coverage
Pilot (v1): two Merge-tier (L3) jobs in `.buildkite/cuda/test-merge.yml` upload
per-model, per-entry-mode coverage as Buildkite artifacts — "Diffusion · Bagel
Test" and "TTS · Qwen3-TTS Base Test". Both run on a single GPU so the pilot is
cheap to reproduce. This is a pilot, not a rollout: no dashboard/visualization
lives in this repo, and nothing is gated on coverage.
**Naming convention**: `coverage-<model_id>-<mode>-<step_id>.xml.gz`, where
`<model_id>` is the model's directory name under
`vllm_omni/diffusion/models/<model_id>/` or
`vllm_omni/model_executor/models/<model_id>/` (e.g. `z_image`, `qwen3_omni`),
`<mode>` is `online` or `offline`, and `<step_id>` is `BUILDKITE_STEP_ID`
(`local` outside Buildkite). `gunzip` before feeding a report to a Cobertura
consumer.
Reports are gzipped because `--cov=vllm_omni` sets coverage's `source`, which
makes it walk the package and emit one `<line>` element per statement for every
file — including the ones a single model's tests never import. That inventory,
not the coverage, is the size: a Bagel report is 7.1 MB, of which 96.5% is
`<line>` elements and 77% belongs to the 840 of 1169 files that run never
imported. A run covering zero lines produces the same 7 MB. gzip takes it to
roughly 410 KB and changes nothing about the data.
`<step_id>` is what keeps two steps that cover the same model apart. Artifacts are
scoped to the build, and a nightly build carries the ready, merge and nightly tiers
at once, so several same-model steps land in one namespace — `qwen3_tts` already has
both a Base and a CustomVoice job, and `minicpmo_4_5` a base and a duplex one. Two
uploads on one path do not fail loudly: an exact-name download can report the path
as ambiguous, and a glob download fetches every match concurrently and renames each
onto the same destination, so whichever finishes last wins. The step id remains the
logical step identifier across retries, while Buildkite's default artifact lookup
selects only the latest attempt — `BUILDKITE_JOB_ID` would instead give every
attempt its own filename, leaving a failed attempt's partial report to be picked up
alongside the real one. A step using `parallelism`/matrix would need
`BUILDKITE_PARALLEL_JOB` added, since its jobs share a step id; neither pilot is
one.
**Opting in a new model**: replace the job's combined `pytest` command with
[`run_cov_split.sh`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/scripts/run_cov_split.sh),
which runs each mode, names the reports, uploads them, and fails the step if
either half or the upload failed:
```yaml
commands:
- |
.buildkite/common/scripts/run_cov_split.sh \
--model-id <model_id> \
--offline <offline test path> \
--online <online test path> \
-- <the job's existing pytest flags, e.g. -m '...' --run-level '...'>
```
Everything after `--` is forwarded verbatim to both runs with quoting preserved,
so an expression like `-m 'advanced_model and cuda'` survives intact and a job can
pass anything else it needs. `--offline`/`--online` are repeatable and only one is
required, so a job with tests in a single mode runs just that half. Runtime is
bounded solely by the step's `timeout_in_minutes` — raise it, since the split
loads the model once per mode.
Before splitting, confirm each half actually collects at least one test under the job's forwarded filters — pytest exits with code 5 ("no tests collected") if a half is empty, which the script reports as a failure, red-ing a job that used to pass as one combined invocation.
Also check the tests' `num_cards` against the job's `mirror_hardwares` preset.
`cuda_marks` turns `num_cards > 1` into `skipif(device_count() < num_cards)`, so a
test asking for more GPUs than the preset provides is silently skipped and its
coverage artifact comes back empty. Splitting such a job produces a file that
looks fine and measures nothing — "Diffusion · Z Image Test" is the current
example (online tests want 4 cards, the job runs on `l4_1`), which is why it is
not a coverage pilot.
Online-mode coverage depends on the `[tool.coverage.run]` `patch`/`sigterm`
settings in `pyproject.toml`: online tests launch the server with
`subprocess.Popen` and stop it with SIGTERM, and without those settings the XML
reflects only the pytest parent process, not the server's code paths.
The upload depends on `buildkite-agent` being callable inside the container. The
`kubernetes` presets (`h100_*`, `*_npu_*`) provide it; the `docker` ones (`l4_*`)
only do because they set `mount-buildkite-agent: true`. A new docker preset that
runs a coverage job needs the same.
List both `run_cov_split.sh` and `pyproject.toml` in every opted-in job's
`source_file_dependencies` — both change what the job measures, so without them a
change there is filtered out of normal PR builds and only surfaces in a later
nightly. `tests/buildkite/test_upload_pipeline.py` asserts this for the pilots.
Editing only the surrounding CI YAML still does not schedule them, so a PR that
touches just the wiring needs a full E2E run (or the commands run on a GPU host)
to produce artifacts. When checking a new model's
artifacts, compare `lines-covered` between the online and offline XML rather than
just confirming both files exist.
### Validation checklist
| Check | Command / location |
+1 -1
View File
@@ -406,7 +406,7 @@ the source for configuration, compatibility, and implementation details.
stage contract.
* **Quantization and memory efficiency:** [Quantization](feature/quantization.md)
resolves per-pipeline or per-component quantization configurations, while
[Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
[Distributed Layerwise Offload](feature/offloader/distributed_layerwise_offload.md)
streams diffusion blocks from host memory within the existing parallel
topology.
@@ -0,0 +1,89 @@
# Diffusion Attention Backend Selection
This document defines the selection and extension contract for diffusion
attention backends. User-facing backend choices, installation, and tuning are
in the [attention backend guides](../../user_guide/diffusion/attention_backends.md).
## Scope
The contract applies to DiT and other diffusion attention layers. It is
separate from vLLM's autoregressive attention selector.
The selection path has four responsibilities:
1. normalize user configuration into `AttentionConfig` and `AttentionSpec`;
2. resolve a spec for an attention role;
3. ask the active platform to validate an explicit backend or choose a
hardware default; and
4. load the selected `AttentionBackend` class from the registry path.
## Resolution contract
`get_attn_backend_for_role()` resolves configuration in this order:
1. exact `per_role[role]`;
2. category `per_role[role_category]`;
3. `default`;
4. platform default.
An explicit resolution returns both the backend class and its
`AttentionSpec`. A platform-default resolution returns the class and `None`.
Layers must therefore treat the spec as optional and must not infer that a
missing spec means SDPA.
Class resolution is cached by backend name, head size, and whether TRTLLM may
be selected automatically. Logging is cached separately by role and source so
the startup record identifies why each role received its backend.
## Registry and platform boundary
`DiffusionAttentionBackendEnum` maps stable configuration names to default
qualified class paths. `register_diffusion_backend()` may replace a path at
runtime without changing the public enum value.
The active `OmniPlatform` owns compatibility policy through
`get_diffusion_attn_backend_cls()`. It must:
- validate explicit selections and fail with an actionable error when the
requested kernel cannot run;
- choose only an available, compatible default when no backend is explicit;
- account for head size and other platform-visible constraints; and
- return a qualified class path implementing `AttentionBackend`.
The selector must not duplicate device capability or package-availability
policy that belongs to the platform.
## Typed backend options
Backend-specific settings remain typed fields on `AttentionSpec` rather than
unstructured keyword dictionaries:
- `quant` is consumed by FlashInfer and TRTLLM with backend-specific value
validation;
- `skip_softmax` is consumed by TRTLLM; and
- `block_sparse` is consumed by block-sparse backends such as RainFusion.
A backend reads only the fields it owns and rejects incompatible values. New
options should be added to a shared typed spec only when more than one backend
shares their semantics; otherwise add a dedicated typed block.
## Model integration contract
Each diffusion attention call declares a stable role and, when appropriate, a
broader category. Model-specific roles permit precise overrides; categories
preserve common `self` or `cross` policies. A model also decides whether its
path is compatible with automatic TRTLLM selection, because only the model
knows whether masking and packing satisfy the kernel contract.
## Adding or changing a backend
An implementation change is complete only when it:
1. implements the `AttentionBackend` interface and exposes a stable enum name;
2. adds platform validation and default routing where appropriate;
3. defines and validates any typed configuration;
4. tests explicit selection, default selection, and incompatible paths;
5. documents installation, fallback behavior, and quality implications in the
corresponding user guide; and
6. updates the overview matrix without moving algorithm details back into the
landing page.
+80
View File
@@ -0,0 +1,80 @@
# CPU Offloading
This document defines the shared architecture for diffusion CPU offloading.
The three strategies have separate user and design pages:
| Strategy | User guide | Design contract |
| --- | --- | --- |
| Model-level | [Guide](../../../user_guide/diffusion/offloader/cpu_offload.md) | [Design](cpu_offload.md) |
| Layerwise | [Guide](../../../user_guide/diffusion/offloader/layerwise_offload.md) | [Design](layerwise_offload.md) |
| Distributed layerwise | [Guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md) | [Design](distributed_layerwise_offload.md) |
## Strategy selection
`OffloadConfig.from_od_config()` converts diffusion configuration into one
`OffloadStrategy`. If callers set multiple flags, selection is deterministic:
1. `DISTRIBUTED_LAYERWISE`;
2. `LAYERWISE`;
3. `MODEL_LEVEL`; and
4. `NONE`.
The factory logs when a higher-priority strategy disables another. It derives
parallel and HSDP state from `DiffusionParallelConfig`; callers do not provide
a separate offload group size.
`get_offload_backend()` then validates platform offload support, resolves the
device, and creates exactly one backend. Returning `None` means offloading is
disabled or unsupported and must not leave partially installed hooks.
## Shared lifecycle
Every backend implements `OffloadBackend`:
- `enable(pipeline)` discovers modules, establishes initial residency, and
installs hooks;
- `disable()` removes owned hooks and resources; and
- `is_enabled()` reports lifecycle state.
`disable()` does not promise to restore the pipeline's original device
placement. The caller owns any subsequent rematerialization.
Hooks are registered through `HookRegistry` and `ModelHook`; offload backends
must use distinct hook names and remove only hooks they own.
## Discovery and topology
Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
- `_dit_modules`;
- `_encoder_modules`;
- `_vae_modules`; and
- `_resident_modules`.
Dotted paths are supported. Legacy pipelines may use the fallback scan of
well-known attribute names, but new integrations should declare components
explicitly.
Ordinary layerwise offload obtains topology from declared
`_layerwise_offload_blocks_attrs` (with a deprecated singular-name
compatibility path). Distributed layerwise offload may instead use a pipeline
`OffloadPlan`, then falls back to declared block attributes and its documented
compatibility discovery. Discovery metadata describes structure only; the
backend remains responsible for transfer, synchronization, and storage
ownership.
## Cross-strategy invariants
- At most one strategy owns offload hooks for a pipeline.
- A parameter has one authoritative host representation while offloaded.
- Device storage is not freed until dependent compute or transfer work is
complete.
- Non-persistent and model-specific buffers remain correct after movement or
checkpoint rematerialization.
- Unsupported parallel or loading combinations fail before hooks mutate the
model.
- Platform streams, events, synchronization, and cache management go through
the vLLM-Omni platform abstraction.
The diffusion [Offloader module design](../../module/diffusion/offloader.md)
describes how these feature contracts fit into the larger diffusion runtime.
@@ -0,0 +1,56 @@
# Model-Level Offload
Model-level offload enforces mutual exclusion between pipeline component
groups. For commands and model support, see the
[model-level user guide](../../../user_guide/diffusion/offloader/cpu_offload.md).
## Generic hook path
`ModelLevelOffloadBackend.enable()` discovers DiTs, encoders, and VAEs. It
makes encoders and VAEs device resident for initialization, then installs
`SequentialOffloadHook` instances:
- each DiT hook offloads all encoders and other DiTs before loading itself;
- each encoder hook offloads every DiT before loading itself; and
- VAEs are not part of the mutual-exclusion hook set.
The hook moves parameters and buffers without recursively calling
`module.to()`. This avoids recursion through modules that retain references to
other transformer blocks. Host tensors are pinned when configured; DTensors
are not pinned through the ordinary tensor path.
Before forward returns control to the model, the hook synchronizes the active
platform so compute cannot race the transfer. XPU uses blocking device-to-host
copies because its allocator cannot safely evict storage while an asynchronous
copy is pending.
## Split-model delegation
If a pipeline provides callable `enable_omni_model_cpu_offload`, the backend
delegates instead of installing generic encoder/DiT hooks. This supports models
whose mutually exclusive components are nested inside one transformer, such
as a reasoner and generator.
The delegated pipeline owns its internal contexts, but must preserve the same
contract:
- only the executing component is device resident;
- enable and disable are idempotent;
- transfers use the supplied device, pinning, and HSDP settings; and
- disabling removes all model-local offload hooks or contexts.
The delegation is currently duck typed. If another independent implementation
is added, formalize the enable/disable pair as a protocol rather than expanding
attribute-name heuristics.
## Extension invariants
- New ordinary pipelines declare both DiT and encoder groups through
`SupportsComponentDiscovery`.
- Models with multiple DiTs offload the inactive DiTs as well as encoders.
- Empty or parameterless modules do not trigger transfers.
- Hook removal does not move modules; callers must not assume original
residency is restored.
Shared selection and lifecycle behavior is defined in the
[CPU Offloading design](README.md).
@@ -7,7 +7,7 @@ either shard those host-side weights across an existing parallel group or keep
the standard loader's rank-local weights and avoid an additional collective.
For user-facing commands, see the
[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md)
[distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
## Status
@@ -56,7 +56,7 @@ AllGather: [full N+1] [full N+2]
Buffers: [current slot] [prefetch slot] [current slot]
```
![DLO double-buffer prefetch pipeline](../figures/dlo/dlo_pipeline.gif)
![DLO double-buffer prefetch pipeline](../../figures/dlo/dlo_pipeline.gif)
The backend uses two shared device buffers, so accelerator weight residency is
bounded by the largest streamed blocks rather than the complete model.
@@ -0,0 +1,57 @@
# Layerwise Offload
Layerwise offload streams transformer blocks through one accelerator. For
commands and model support, see the
[layerwise user guide](../../../user_guide/diffusion/offloader/layerwise_offload.md).
## Hook ring
`LayerWiseOffloadBackend` discovers every streamable block and registers one
`LayerwiseOffloadHook` per block. Hooks form a ring: each hook owns the pinned
host representation of its registered block and a reference to the next block
it must prefetch; the final hook prefetches the first block for the next
denoising iteration.
During initialization, parameters and buffers are grouped by dtype and copied
into contiguous pinned CPU tensors. Original tensor storage is replaced with
an empty or meta placeholder while identity and module structure are retained.
Metadata records names, offsets, shapes, and dtypes for rematerialization.
## Forward lifecycle
For each block:
1. the previous hook's completion event ensures this block is materialized;
2. the copy stream asynchronously allocates and fills the next block's device
tensors;
3. parameter and buffer storage is rebound to views of those tensors;
4. the current block computes on the normal stream; and
5. its post-forward hook replaces device storage with placeholders after all
dependent work is safe.
The copy stream waits on the compute stream before reusing storage, and the
compute stream waits on the per-block prefetch event before forward. These
dependencies, not incidental global synchronization, define correctness.
## Topology contract
A DiT declares ordered block-container attributes using
`_layerwise_offload_blocks_attrs` or an `OffloadPlan`. Each resolved item must
be an executable `nn.Module`; ordering must match forward execution. Multiple
containers are concatenated in execution order.
Non-block DiT modules, encoders, VAEs, and declared resident modules are not
streamed by this backend.
## Invariants and limitations
- Each offloaded tensor has one pinned host backing store.
- Parameter identity is preserved when storage is rebound.
- Blocks with heterogeneous shapes and dtypes use their recorded metadata;
implementations must not assume identical blocks.
- The ring is initialized before the first forward so block zero is available.
- This backend is single-device; distributed weight reconstruction belongs to
[Distributed Layerwise Offload](distributed_layerwise_offload.md).
Shared selection and lifecycle behavior is defined in the
[CPU Offloading design](README.md).
+1 -1
View File
@@ -186,7 +186,7 @@ Distributed layerwise offload has an additional boundary: its sharded mmap
AllGather path is not compatible with online quantization because the path
expects already materialized checkpoint weights. Use the rank-local
`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when
combining these features. See the [DLO design document](distributed_layerwise_offload.md)
combining these features. See the [DLO design document](offloader/distributed_layerwise_offload.md)
for the full compatibility matrix.
## Adding a quantization backend
+4 -2
View File
@@ -2,8 +2,10 @@
Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. Usage — the config keys
and how to pick an operating point — is in
[Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). This page explains
the algorithm.
[TRTLLM Attention](../../user_guide/diffusion/attention_backends/trtllm.md#skip-softmax).
The shared selector contract is documented in
[Diffusion Attention Backend Selection](attention_backend_selection.md). This
page explains the algorithm.
## Motivation
+10 -3
View File
@@ -53,15 +53,22 @@ implementation contract; it is not, by itself, a general support claim.
#### 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:
guides list selectable backends, platform defaults, installation, and tuning.
The design contracts separate selection mechanics from backend algorithms:
- [Attention Backend Selection](feature/attention_backend_selection.md)
- [Skip-Softmax](feature/skip_softmax.md)
#### CPU offloading
- [Overview and Shared Contracts](feature/offloader/README.md)
- [Model-Level Offload](feature/offloader/cpu_offload.md)
- [Layerwise Offload](feature/offloader/layerwise_offload.md)
- [Distributed Layerwise Offload](feature/offloader/distributed_layerwise_offload.md)
- [Cache-DiT](feature/cache_dit.md)
- [TeaCache](feature/teacache.md)
- [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
## Infrastructure and Performance
+1 -1
View File
@@ -1,6 +1,6 @@
# Prometheus Metrics Design
This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics. This iteration covers the text and audio surface; diffusion / image / video families land in a follow-up.
This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics.
## Objectives
+2 -2
View File
@@ -48,10 +48,10 @@ defines the shared configuration and backend extension points.
| 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) |
| Move weights between host and device memory | [CPU Offloading](../user_guide/diffusion/cpu_offload.md) | [CPU Offloading](../design/feature/offloader/README.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) |
| Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Attention Backend Selection](../design/feature/attention_backend_selection.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 |
+3 -1
View File
@@ -42,7 +42,8 @@ th {
| `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` | ✅︎ | ✅︎ | | | — |
| `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — |
| `LTX2DistilledPipeline` | LTX-2 / LTX-2.3 full-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + `Lightricks/LTX-2.3` upsampler | ✅︎ | ✅︎ | | | — |
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -58,6 +59,7 @@ th {
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
| `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+3 -3
View File
@@ -10,12 +10,12 @@ A: Refer to [GPU memory calculation and configuration](../configuration/gpu_memo
> Q: I encounter some bugs or CI problems, which is urgent. How can I solve it?
A: At first, you can check current [issues](https://github.com/vllm-project/vllm-omni/issues) to find possible solutions. If none of these satisfy your demand and it is urgent, please find these [volunteers](https://docs.vllm.ai/projects/vllm-omni/en/latest/community/volunteers/) for help.
A: Search the current [issues](https://github.com/vllm-project/vllm-omni/issues) first. If no existing issue matches a bug or CI failure, open a [bug report](https://github.com/vllm-project/vllm-omni/issues/new?template=400-bug-report.yml) with a minimal reproducer, environment details, and relevant logs. For technical or usage questions, use the `sig-omni` channel in [vLLM Slack](https://slack.vllm.ai/) or the [vLLM Forum](https://discuss.vllm.ai/).
> Q: Does vLLM-Omni support AWQ or any other quantization?
A: We plan to introduce GGUF FP8 prequantized models and online FP8 quantization in version 0.16.0. Support for other quantization types will follow in future releases. For details, please see our [Q1 quantization roadmap](https://github.com/vllm-project/vllm-omni/issues/1057).
A: AWQ is available through the inherited vLLM quantization registry, and vLLM-Omni supports several additional methods. Actual support and validation vary by model, component, hardware, and checkpoint. See the [quantization overview](../user_guide/quantization/overview.md) for validated combinations and method-specific guides.
> Q: Does vLLM-Omni support multimodal streaming input and output?
A: Not yet. We already put it on the [Roadmap](https://github.com/vllm-project/vllm-omni/issues/165). Please stay tuned!
A: Yes, for specific models and serving paths. Current capabilities include [streaming video input](../serving/video_stream_api.md), [streaming diffusion output](../user_guide/diffusion/execution_modes.md#streaming-output), and experimental full-duplex real-time audio described in the [project overview](../README.md). Check the model and API documentation because input and output support is model-specific.
+50 -26
View File
@@ -1,23 +1,23 @@
# Production Metrics
vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server. This page covers the text and audio surface; diffusion / image / video metrics are tracked in a follow-up PR.
vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server.
```bash
vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8000 --log-stats
curl http://localhost:8000/metrics
```
**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm:omni_*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm_omni:*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
## Metric Namespaces
| Prefix | Source | Present when |
|--------|--------|--------------|
| `vllm:omni_` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
| `vllm_omni:` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
| `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage |
| `http_` / `process_` | Uvicorn / Python runtime | Always |
## Pipeline-Level Metrics (`vllm:omni_`)
## Pipeline-Level Metrics (`vllm_omni:`)
Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the full multi-stage pipeline.
@@ -25,42 +25,66 @@ Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the
| Metric | Type | Labels | Description |
|--------|------|--------|-------------|
| `vllm:omni_num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
| `vllm_omni:num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
| `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
| `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
### Latency
| Metric | Type | Labels | Description |
|--------|------|--------|-------------|
| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
| `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
## Audio Modality Metrics (`vllm:omni_`)
## Audio Modality Metrics (`vllm_omni:`)
Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label.
| Metric | Type | Extra label | Description |
|--------|------|-------------|-------------|
| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
| `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
| `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
| `vllm_omni:audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
| `vllm_omni:audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
| `vllm_omni:audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
| `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
| `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition.
## Cross-Stage Transfer Metrics (`vllm:omni_`)
## Diffusion Metrics (`vllm_omni:`)
Per-request timing breakdowns for diffusion (image/video) stages. Emitted at request finalize when `stage_metrics.diffusion_metrics` is present. All carry `{model_name, stage, replica}`.
| Metric | Type | Description |
|--------|------|-------------|
| `vllm_omni:diffusion_exec_s` | Histogram | DiT forward pass execution time per request in seconds |
| `vllm_omni:diffusion_exec_per_step_s` | Histogram | DiT forward pass execution time per denoising step in seconds (`exec_s / num_inference_steps`) |
| `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time per request in seconds |
| `vllm_omni:diffusion_postprocess_s` | Histogram | Diffusion output postprocessing (VAE decode) time per request in seconds |
`diffusion_exec_s` uses `SECONDS_BUCKETS` (50 ms → 300 s); the other three use `SECONDS_FAST_BUCKETS` (1 ms → 60 s). `diffusion_exec_per_step_s` is skipped when `num_inference_steps` is not set.
### Example PromQL
```promql
# P99 DiT exec time by stage
histogram_quantile(0.99, rate(vllm_omni:diffusion_exec_s_bucket[5m]))
# Ratio of preprocess to total exec wall time
rate(vllm_omni:diffusion_preprocess_s_sum[5m])
/ rate(vllm_omni:diffusion_exec_s_sum[5m])
```
## Cross-Stage Transfer Metrics (`vllm_omni:`)
Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed.
| Metric | Type | Description |
|--------|------|-------------|
| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
| `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
| `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
| `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
| `vllm_omni:transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
## vLLM Engine Metrics (`vllm:`)
@@ -83,9 +107,9 @@ For the full list of upstream metrics, see [the vLLM docs](https://github.com/vl
| Metric group | Multi-stage LLM (Qwen3-Omni) |
|---|---|
| `vllm:omni_` request tracking + latency | With `--log-stats` |
| `vllm:omni_` audio modality | With `--log-stats`, if pipeline has a talker stage |
| `vllm:omni_` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
| `vllm_omni:` request tracking + latency | With `--log-stats` |
| `vllm_omni:` audio modality | With `--log-stats`, if pipeline has a talker stage |
| `vllm_omni:` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
| `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` |
| `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` |
@@ -94,5 +118,5 @@ For the full list of upstream metrics, see [the vLLM docs](https://github.com/vl
- All time-bearing metrics use the `_s` suffix (values in seconds). Buckets are `SECONDS_BUCKETS` for e2e / generation-style values and `SECONDS_FAST_BUCKETS` (1 ms → 60 s) for the fine-grained transfer and audio-underrun values.
- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
- Sizes use the `_bytes` suffix.
- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm:omni_audio_ttfp_s` for audio) rather than a single metric with a `modality` label.
- All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm_omni:`.
- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm_omni:audio_ttfp_s` for audio) rather than a single metric with a `modality` label.
+111 -413
View File
@@ -1,494 +1,192 @@
# Diffusion Attention Backends
This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, the per-platform defaults, and how to use SageAttention.
Use this page to select and configure a diffusion attention backend. Backend-
specific installation, tuning, and compatibility details live in separate
guides so that the selection contract stays easy to scan.
## Overview
Diffusion attention backend selection applies to DiT attention in image and
video generation models. It does **not** change autoregressive LLM attention,
which uses vLLM's attention selector.
Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured.
For the internal selector, registry, and platform contract, see
[Attention Backend Selection](../../design/feature/attention_backend_selection.md).
This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. It does **not** affect autoregressive (LLM) attention paths — those go through vLLM's own attention backend selector.
## Choose a guide
The full set of backends and their platform defaults is in the **Backend Options** and **Platform Defaults** sections below. If no attention backend is configured, vLLM-Omni asks the current platform to choose the default.
| Need | Guide |
| --- | --- |
| Select a conservative or platform-native dense kernel | [Dense Backends](attention_backends/dense_backends.md) |
| Use FlashInfer trtllm-gen FMHA, Skip-Softmax, or TRTLLM SAGE quantization | [TRTLLM Attention](attention_backends/trtllm.md) |
| Install and use SageAttention 2.2 or SageAttention3 | [SageAttention](attention_backends/sage.md) |
| Match training or rollout kernels loaded from Hugging Face | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
| Use block-sparse video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
## Backend Options
| Value | Notes |
|---|---|
| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). Dense BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Packed paths can provide `cu_seqlens` directly. Supports optional **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
| `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. |
| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
| `FLASHINFER_ATTN` | Uses FlashInfer's batch-prefill wrapper and supports mixed Q/K and V input dtypes through backend-specific configuration. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer` >= 0.6.16rc1 for mixed-dtype configurations. |
| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
| `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
| `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. |
| `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. |
| `RAINFUSION_ATTN` | MindIE-SD **RainFusion** block-sparse video attention — see [below](#rainfusion_attn-backend-and-block-sparse-video-attention). Ascend NPU only; requires `mindiesd`. Delegates to `FLASH_ATTN` for anything that is not a packed video sequence. |
## Backend options
| Value | Family | Primary use | Detail |
| --- | --- | --- | --- |
| `TORCH_SDPA` | Dense | Conservative reference; always available | [Dense Backends](attention_backends/dense_backends.md#torch_sdpa) |
| `FLASH_ATTN` | Dense | FlashAttention 4/3/2 depending on the installed package and GPU | [Dense Backends](attention_backends/dense_backends.md#flash_attn) |
| `CUDNN_ATTN` | Dense | Mask-heavy DiTs on Blackwell with cuDNN 9.5 or newer | [Dense Backends](attention_backends/dense_backends.md#cudnn_attn) |
| `FLASHINFER_ATTN` | Dense or quantized | FlashInfer batch prefill; optional mixed Q/K and V dtypes | [Dense Backends](attention_backends/dense_backends.md#flashinfer_attn) |
| `TRTLLM_ATTN` | Dense, sparse, or quantized | Datacenter Blackwell with `head_dim=128` and compatible packed paths | [TRTLLM Attention](attention_backends/trtllm.md) |
| `SAGE_ATTN` | Quantized | SageAttention 2.2 INT8 attention | [SageAttention](attention_backends/sage.md#sage_attn) |
| `SAGE_ATTN_3` | Quantized | SageAttention3 on Blackwell | [SageAttention](attention_backends/sage.md#sage_attn_3) |
| `FLASH_ATTN_HUB` | Hub kernel | FlashAttention 2 from Hugging Face `kernels` | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
| `FLASH_ATTN_3_HUB` | Hub kernel | FlashAttention 3 from Hugging Face `kernels` on Hopper or newer | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
| `RAINFUSION_ATTN` | Block sparse | MindIE-SD RainFusion video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
## Configuration
Diffusion attention backends can be configured three ways, in priority order:
Backend selection is resolved in this order:
1. **`--diffusion-attention-config`** — structured per-role config (highest priority).
2. **`--diffusion-attention-backend` / `DIFFUSION_ATTENTION_BACKEND` env var** — global shorthand that sets the default backend.
3. **Platform default** — used when nothing is configured.
1. `--diffusion-attention-config` per-role configuration.
2. `--diffusion-attention-backend` or `DIFFUSION_ATTENTION_BACKEND` as a
global default.
3. The current platform's default.
`--diffusion-attention-backend` is shorthand for `--diffusion-attention-config.default.backend`. It may be combined with `--diffusion-attention-config.per_role.*` overrides, but is mutually exclusive with `--diffusion-attention-config.default.backend`.
`--diffusion-attention-backend` is shorthand for
`--diffusion-attention-config.default.backend`. Do not pass it together with an
explicit `default.backend` in the structured configuration.
### Global default
Set the default backend for every diffusion attention layer:
```bash
# CLI flag
vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN
# Environment variable (also recognized for backwards compatibility)
export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN
# Backwards-compatible environment variable
export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
```
### Per-role configuration
Roles are free-form strings declared by each diffusion model. The two common categories are `"self"` and `"cross"`; model-specific roles (e.g. `"ltx2.audio_to_video"`) may also be declared. A role string is matched in this order:
Roles are declared by each diffusion model. Common categories are `self` and
`cross`; a model may also use a more specific role such as
`ltx2.audio_to_video`. Resolution order is:
1. Exact `per_role[role]` match
2. `per_role[role_category]` fallback (e.g. `"ltx2.audio_to_video"` → `"cross"`)
3. `default`
4. Platform default
Use vLLM-style dotted flags or one JSON blob:
1. Exact `per_role[role]` match.
2. Category `per_role[role_category]` match.
3. `default`.
4. Platform default.
```bash
# Dotted flags
vllm-omni serve <model> \
--diffusion-attention-config.default.backend FLASH_ATTN \
--diffusion-attention-config.per_role.cross.backend TORCH_SDPA
--diffusion-attention-config.default.backend FLASH_ATTN \
--diffusion-attention-config.per_role.cross.backend TORCH_SDPA
# JSON
# Equivalent JSON
vllm-omni serve <model> \
--diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
--diffusion-attention-config \
'{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
```
A backend that needs configuration exposes it as a typed field on the spec:
`TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)) and
`RAINFUSION_ATTN`'s RainFusion block (see [below](#rainfusion_attn-backend-and-block-sparse-video-attention)):
```bash
--diffusion-attention-config.default.backend TRTLLM_ATTN \
--diffusion-attention-config.default.skip_softmax.target_sparsity 0.5
```
### Programmatic API
When constructing `OmniDiffusionConfig` directly:
### Python API
```python
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec, OmniDiffusionConfig
from vllm_omni.diffusion.data import (
AttentionConfig,
AttentionSpec,
OmniDiffusionConfig,
)
config = OmniDiffusionConfig(
diffusion_attention_config=AttentionConfig(
default=AttentionSpec(backend="FLASH_ATTN"),
per_role={
"cross": AttentionSpec(backend="TORCH_SDPA"),
},
per_role={"cross": AttentionSpec(backend="TORCH_SDPA")},
),
...,
)
```
A plain dict is also accepted and normalized to `AttentionConfig`.
Backend-specific typed blocks are documented with their consumers:
#### FlashInfer QK16/V8 quantized attention
- `quant`: [FlashInfer](attention_backends/dense_backends.md#flashinfer-quantized-attention)
and [TRTLLM SAGE](attention_backends/trtllm.md#sage-quantization).
- `skip_softmax`: [TRTLLM Skip-Softmax](attention_backends/trtllm.md#skip-softmax).
- `block_sparse`: [RainFusion](attention_backends/rainfusion.md#configuration).
`FLASHINFER_ATTN` accepts quantization options through `AttentionSpec.quant`. For QK16/V8:
```python
config = OmniDiffusionConfig(
diffusion_attention_config=AttentionConfig(
default=AttentionSpec(
backend="FLASHINFER_ATTN",
quant=AttnQuantSpec(
dtype_qk="bfloat16",
dtype_vo="fp8_e4m3",
),
),
),
...,
)
```
`dtype_qk` configures the Q and K input dtypes. `dtype_vo` configures the V input dtype.
## Platform Defaults
## Platform defaults
### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
Auto-route preference, in order:
The CUDA auto-route preference is:
1. `TRTLLM_ATTN` — on **datacenter** Blackwell (sm_100 / sm_103) when `flashinfer` is installed, the model's `head_dim` is 128, **and the model declares a compatible packed/mask-free path**
2. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
3. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
4. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
5. `TORCH_SDPA` — last resort
1. `TRTLLM_ATTN` on datacenter Blackwell (sm_100/sm_103) when FlashInfer is
available, `head_dim=128`, and the model declares a compatible packed or
mask-free path.
2. `CUDNN_ATTN` when cuDNN 9.5 or newer is available.
3. `FLASHINFER_ATTN` when FlashInfer is available but cuDNN is too old.
4. `FLASH_ATTN` when a compatible package is installed.
5. `TORCH_SDPA`.
`TRTLLM_ATTN` is skipped on workstation Blackwell (sm_120 / sm_121) and for any `head_dim != 128`, so those GPUs keep the `CUDNN_ATTN` route described below.
`TRTLLM_ATTN` is not auto-selected on workstation Blackwell
(sm_120/sm_121), for other head dimensions, or for paths that require an
unsupported mask.
`TRTLLM_ATTN` outranks `CUDNN_ATTN` on datacenter Blackwell for compatible packed/mask-free pipelines. Workstation Blackwell (sm_120 / sm_121) and pipelines that require attention masks retain their normal fallback.
### Hopper, Ada, and Ampere
The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.
The CUDA auto-route uses `FLASH_ATTN` when available and otherwise falls back
to `TORCH_SDPA`. `CUDNN_ATTN` and `FLASHINFER_ATTN` remain explicit options.
**Why CUDNN_ATTN by default**: on mask-heavy diffusion models (HunyuanVideo-1.5, Qwen-Image), cuDNN's pinned FMHA kernel sidesteps a PyTorch SDPA dispatch quirk where the unpinned dispatcher picks `EFFICIENT_ATTENTION` (~25 ms) for masked calls instead of cuDNN (~11 ms). The pin gives 2× e2e on HV-1.5 with no regression on lighter models.
Other platforms validate an explicit backend and choose their own default
through the platform implementation. Check the startup log to confirm the
resolved backend.
### Hopper (sm_90) / Ada (sm_89) / Ampere (sm_80–sm_86)
## Choosing a backend manually
Auto-route preference:
Override the platform default when you need:
1. `FLASH_ATTN` — when `flash-attn` is installed
2. `TORCH_SDPA` — fallback
- a correctness reference (`TORCH_SDPA`);
- a backend-specific workaround;
- training/rollout kernel alignment (Hub backends); or
- an explicitly validated sparse or quantized speedup.
`CUDNN_ATTN` and `FLASHINFER_ATTN` are still selectable via env var on these GPUs but are not in the auto-route — FlashAttention 2 is the well-tuned path on pre-Blackwell hardware.
The startup log prints the resolved backend and whether it came from explicit
configuration or platform defaulting. If no resolution message appears, check
earlier logs for diffusion-stage initialization failures.
## TRTLLM_ATTN Backend and Skip-Softmax
## Reference benchmark
`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA. Dense it is a BF16 backend like the others; on top of
that it can enable **Skip-Softmax**, a sparse-attention mode that trades a little fidelity for speed
(algorithm: [design doc](../../design/feature/skip_softmax.md)).
The following BF16 results were measured on an sm_120 RTX Pro 6000 Blackwell
with the same prompt and seed across runs. Treat them as reference results, not
portable guarantees.
Enable it through the typed `skip_softmax` block on the attention spec:
| Key | Valid values | Meaning |
|---|---|---|
| `target_sparsity` | finite, `[0, 1]` | Operating point on a calibrated curve. Needs a calibration for the model. |
| `threshold` | finite, `≥ 0` | Direct threshold, no calibration needed. Mutually exclusive with `target_sparsity`. |
| `disabled_until_timestep` | finite, `[0, 1]` | Holds early, high-noise steps dense; skip turns on once the normalized timestep `t ≤ D`. |
```bash
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
"skip_softmax": {"target_sparsity": 0.65, "disabled_until_timestep": 0.86}}}'
```
Programmatically the same block is a typed `SkipSoftmaxSpec` (values validated at construction):
```python
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
AttentionConfig(
default=AttentionSpec(
backend="TRTLLM_ATTN",
skip_softmax=SkipSoftmaxSpec(target_sparsity=0.65, disabled_until_timestep=0.86),
),
)
```
**Start at `target_sparsity=0.65, disabled_until_timestep=0.86`** and raise `target_sparsity` only
as far as your quality bar allows — the output diverges from dense as it climbs. The gain grows with
sequence length, since Skip-Softmax only accelerates attention. Dense BF16 is `TRTLLM_ATTN` with no
`skip_softmax` block.
Requires datacenter Blackwell with `head_dim == 128`; elsewhere, or without FlashInfer, selecting
`TRTLLM_ATTN` raises rather than silently degrading.
## TRTLLM_ATTN SAGE Quantization
On top of dense/Skip-Softmax, `TRTLLM_ATTN` can run **SAGE** (SageAttention) — the trtllm-gen
FMHA kernel with per-block quantized Q/K and per-channel fp8 V, so the two attention matmuls run
in low precision. Enable it through the typed `quant` block on the attention spec:
| Key | Valid values | Meaning |
|---|---|---|
| `dtype_qk` | `int8`, `fp8_e4m3` | Quantization dtype for Q and K. Absent ⇒ dense (SAGE off). |
| `q_block_size` | `1`, `4`, `16` | Per-token block size for Q scales (default `1`). |
| `k_block_size` | `1`, `4`, `16` | Per-token block size for K scales (default `16`). |
V is always quantized per-channel to fp8_e4m3 and K-smoothing is applied inside the routine, so
neither is a user knob. Only block sizes `1`, `4`, `16` have compiled kernels.
The `quant` block is a shared `AttnQuantSpec` consumed by more than one backend; each reads the
fields that apply and rejects the rest. `TRTLLM_ATTN` (SAGE) reads `dtype_qk` (`int8` / `fp8_e4m3`),
`q_block_size`, and `k_block_size`; `FLASHINFER_ATTN` reads `dtype_qk` (`float16` / `bfloat16`),
`dtype_vo`, and `flashinfer_backend` (see its section). A `dtype_qk` from the wrong set raises.
```bash
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
"quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}}}'
```
```python
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
AttentionConfig(
default=AttentionSpec(
backend="TRTLLM_ATTN",
quant=AttnQuantSpec(dtype_qk="fp8_e4m3", q_block_size=1, k_block_size=16),
),
)
```
**Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel
availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103**
(B300); `int8` QK kernels are compiled for **SM100 only**.
## RAINFUSION_ATTN Backend and Block-Sparse Video Attention
`RAINFUSION_ATTN` runs MindIE-SD's RainFusion (`rf_v2`) kernel on Ascend NPU. It pools each
128-token block of keys, ranks them per query block, and computes attention against only the
top-scoring ones. Video tokens are rearranged into `(t, h, w)` order first, so the kept blocks are
spatiotemporal neighbours rather than arbitrary rows of the packed sequence.
Only the video segment is sparsified. The prefix rows (text, visual conditions, audio) and the
first-frame blocks are always kept dense, which is why the realized sparsity is lower than the
nominal `sparsity` you configure. Anything the kernel cannot handle — a warmup denoise step, an
exempt layer, a sequence with no published video geometry, a video segment under 32 blocks —
delegates to `FLASH_ATTN`, so a model can select this backend unconditionally.
Enable it through the typed `block_sparse` block on the attention spec, shared by every
block-sparse backend:
| Key | Valid values | Meaning |
|---|---|---|
| `sparsity` | finite, `[0, 1]` | Nominal fraction of key blocks dropped per query block. `0` disables sparsity. Defaults to `0.8`. |
| `start_step` | `≥ 0` | Keep the first N denoise steps dense. Layout is decided early, so these steps dominate structural fidelity. |
| `skip_layers` | index selector, e.g. `"0-3,38"` | DiT blocks that always stay dense. |
```bash
vllm-omni serve MiniMaxAI/MiniMax-H3 \
--diffusion-attention-config '{"default": {"backend": "RAINFUSION_ATTN",
"block_sparse": {"sparsity": 0.8, "start_step": 0}}}'
```
Programmatically the same block is a typed `BlockSparseSpec` (values validated at construction):
```python
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec
AttentionConfig(
default=AttentionSpec(
backend="RAINFUSION_ATTN",
block_sparse=BlockSparseSpec(sparsity=0.8, start_step=0, skip_layers="0-1"),
),
)
```
**Tune in the order `start_step` → `sparsity` → `skip_layers`.** Raise `start_step` first: it is the
cheapest way to recover structure, because the early high-noise steps decide global layout while the
later steps only refine texture. Then walk `sparsity` up to your quality bar. Reach for
`skip_layers` last, once an A/B against dense at the same seed points at specific blocks.
Requires Ascend NPU with `mindiesd`; selecting it on any other platform raises. It is also
incompatible with ring sequence parallelism, since `rf_v2` needs the whole key sequence to rank
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.
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.
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.
## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
| Model | Shape | TORCH_SDPA | CUDNN_ATTN | FLASHINFER_ATTN |
|---|---|---|---|---|
| Model | Shape | `TORCH_SDPA` | `CUDNN_ATTN` | `FLASHINFER_ATTN` |
| --- | --- | ---: | ---: | ---: |
| HunyuanVideo-1.5 (T2V) | 480p / 33f / 50 steps | 147.05 s | **73.02 s** | 127.84 s |
| Wan 2.2 14B (T2V) | 480p / 33f / 40 steps | 117.75 s | 117.17 s | **115.07 s** |
| Qwen-Image (T2I) | 1024² / 50 steps | 17.41 s | **15.14 s** | 16.02 s |
| FLUX.2-dev (T2I) | 1024² / 50 steps, TP=2 | 53.62 s | **53.30 s** | 54.94 s |
Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
Mask-heavy DiTs favored `CUDNN_ATTN`; lighter-mask or TP-saturated workloads
were close enough that users should benchmark their exact model and shape.
## Known Limitations
## Compatibility anchors
### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
The following headings preserve links to sections that moved into dedicated
guides.
LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
## TRTLLM_ATTN Backend and Skip-Softmax
**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
See [TRTLLM Attention](attention_backends/trtllm.md#skip-softmax).
```bash
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
--model Lightricks/LTX-2 ...
```
## TRTLLM_ATTN SAGE Quantization
### FlashAttention-4 on Blackwell
See [TRTLLM Attention](attention_backends/trtllm.md#sage-quantization).
Install the optional CUDA 13 extra to use the CuTe-based FA4 path:
## RAINFUSION_ATTN Backend and Block-Sparse Video Attention
```bash
pip install 'vllm-omni[fa4]'
```
`FLASH_ATTN` prefers `flash_attn.cute` on Blackwell and falls back to FA3/FA2
when it is unavailable. Version `4.0.0b18` is required; earlier beta
wheels had known JIT failures on Blackwell.
## Choosing a Backend Manually
### When to override the default
- **Quality validation**: compare a new backend against `TORCH_SDPA` as the reference, since SDPA's default dispatcher is the most extensively tested.
- **Lossy speedup hunting**: try `SAGE_ATTN` (INT8 quantized) on diffusion outputs — typically indistinguishable visually but always validate.
- **Workaround for known issues**: see Known Limitations above.
### Verifying which backend is in use
The startup log prints one of:
```
Using diffusion attention backend 'CUDNN_ATTN' # explicit override
Defaulting to diffusion attention backend CUDNN_ATTN ... # auto-route
Defaulting to diffusion attention backend SDPA # nothing else available
```
If you don't see one of these, the model didn't reach diffusion stage init — check earlier logs for failures.
See [RainFusion](attention_backends/rainfusion.md).
## SageAttention Installation
vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
Build from source:
```bash
git clone https://github.com/thu-ml/SageAttention.git
cd SageAttention
export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
pip install . --no-build-isolation
```
Quick check:
```bash
python -c "import sageattention; print(sageattention.__file__)"
```
See [SageAttention](attention_backends/sage.md#installation).
## SageAttention3 Installation
vLLM-Omni expects SageAttention3 to be installed into the same Python environment as vLLM-Omni.
Build from source:
```bash
git clone https://github.com/thu-ml/SageAttention.git
cd SageAttention/sageattention3_blackwell
python setup.py install
```
Quick check:
```bash
python -c "import sageattn3; print(sageattn3.__file__)"
```
Notes:
- `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class.
- SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness.
See [SageAttention](attention_backends/sage.md#sageattention3-installation).
## HuggingFace Kernels Hub Backends
To achieve perfect numerical consistency between **training** (typically using HuggingFace Diffusers with Hub kernels) and **serving/rollout** (using vLLM-Omni), you can use the Hub-based attention backends. This eliminates numerical drift / sampling divergence caused by executing different local kernel versions during rollout.
The following backend options are supported:
- `FLASH_ATTN_HUB` (HuggingFace `kernels-community/flash-attn2`)
- `FLASH_ATTN_3_HUB` (HuggingFace `kernels-community/flash-attn3`, Hopper sm_90+ only)
### Installation
To use these backends, you must install the `kernels` library:
```bash
pip install kernels==0.14.1
```
If the `kernels` library is not available in the environment, vLLM-Omni will log a warning and fall back gracefully to the corresponding local backend implementations (`FLASH_ATTN`). On CUDA GPUs below Hopper (compute capability < 9.0), `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
### Usage
Select a Hub backend using the global CLI flag or environment variables:
```bash
# Environment variable
export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB
# CLI flag
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN_HUB
```
## Usage Examples
### Default (auto-route)
```bash
python examples/offline_inference/text_to_video/text_to_video.py \
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
--prompt "A dog running across a field of golden wheat." \
--height 480 --width 832 --num-frames 33 \
--num-inference-steps 50 --seed 42 --guidance-scale 6.0 \
--output hv15.mp4
```
On Blackwell this picks `CUDNN_ATTN` automatically. Check the log for the `Defaulting to ...` line.
### Explicit backend selection
```bash
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
--model Lightricks/LTX-2 \
--prompt "A dog running across a field of golden wheat." \
--height 480 --width 832 --num-frames 33 \
--num-inference-steps 40 --seed 42 --guidance-scale 4.0 \
--output ltx2.mp4
```
### SageAttention (lossy)
```bash
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
--prompt "A dog running across a field of golden wheat." \
--height 480 --width 832 --num-frames 33 \
--num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
--tensor-parallel-size 2 \
--output hv15_sage.mp4
```
Example: Wan2.2 TI2V 5B
```bash
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--prompt "A dog running across a field of golden wheat." \
--height 704 --width 1280 --num-frames 49 \
--num-inference-steps 30 --seed 42 --guidance-scale 5.0 \
--tensor-parallel-size 2 \
--output outputs/wan22_sage.mp4
```
### Enable SageAttention3
Example:
```bash
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_to_video/text_to_video.py \
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
--prompt "A dog running across a field of golden wheat." \
--height 480 --width 832 --num-frames 33 \
--num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
--tensor-parallel-size 2 \
--output outputs/hv15_sage3.mp4
```
See [Hugging Face Hub Backends](attention_backends/huggingface_hub.md).
@@ -0,0 +1,110 @@
# Dense Attention Backends
Use these backends when you want dense attention without backend-specific
sparsity. Start with `TORCH_SDPA` for correctness comparisons, then benchmark
the fastest compatible kernel for your model, shape, and hardware.
For selection precedence and per-role configuration, see the
[attention backend overview](../attention_backends.md).
## `TORCH_SDPA`
`TORCH_SDPA` calls PyTorch `scaled_dot_product_attention` and lets PyTorch's
dispatcher choose the implementation. It is always available and is the most
conservative reference when validating another backend.
```bash
vllm-omni serve <model> --diffusion-attention-backend TORCH_SDPA
```
## `FLASH_ATTN`
`FLASH_ATTN` uses the installed FlashAttention implementation. On Blackwell it
prefers the CuTe-based FlashAttention 4 path and falls back to
FlashAttention 3 or 2 when that path is unavailable. On Hopper, Ada, and
Ampere it is the preferred automatic route when a compatible package is
installed.
```bash
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN
```
### FlashAttention 4 on Blackwell
Install the optional CUDA 13 extra:
```bash
pip install 'vllm-omni[fa4]'
```
Version `4.0.0b18` is required; earlier beta wheels had known JIT failures on
Blackwell. If the CuTe path is unavailable, the backend falls back to the
compatible FlashAttention 3 or 2 path.
## `CUDNN_ATTN`
`CUDNN_ATTN` pins PyTorch SDPA to `CUDNN_ATTENTION`. It is particularly useful
for mask-heavy DiTs and is automatically preferred on Blackwell when cuDNN
9.5 or newer is available and the higher-priority TRTLLM route is not
compatible.
```bash
vllm-omni serve <model> --diffusion-attention-backend CUDNN_ATTN
```
### LTX-2.0 limitation
LTX-2 audio attention has a symbolic head dimension during `torch.compile`
tracing. The cuDNN SDPA selector rejects that symbolic dimension and Dynamo
aborts compilation. This is tracked in
[issue #3121](https://github.com/vllm-project/vllm-omni/issues/3121).
Use `FLASHINFER_ATTN` or `TORCH_SDPA` as a workaround:
```bash
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN \
python examples/offline_inference/text_to_video/text_to_video.py \
--model Lightricks/LTX-2 ...
```
## `FLASHINFER_ATTN`
`FLASHINFER_ATTN` uses FlashInfer's batch-prefill wrapper. It is an explicit
option on CUDA platforms and an automatic Blackwell fallback when FlashInfer
is installed but cuDNN is too old for `CUDNN_ATTN`.
```bash
vllm-omni serve <model> --diffusion-attention-backend FLASHINFER_ATTN
```
### FlashInfer quantized attention
The backend accepts an `AttentionSpec.quant` block. For QK16/V8, keep Q and K
in FP16 or BF16 and use FP8 E4M3 for V:
```python
from vllm_omni.diffusion.data import (
AttentionConfig,
AttentionSpec,
AttnQuantSpec,
OmniDiffusionConfig,
)
config = OmniDiffusionConfig(
diffusion_attention_config=AttentionConfig(
default=AttentionSpec(
backend="FLASHINFER_ATTN",
quant=AttnQuantSpec(
dtype_qk="bfloat16",
dtype_vo="fp8_e4m3",
),
),
),
...,
)
```
`dtype_qk` controls Q and K; `dtype_vo` controls V. Mixed-dtype
configurations require FlashInfer 0.6.16rc1 or newer. The shared quantization
schema is also consumed by TRTLLM, but each backend validates its own allowed
fields and values; see [TRTLLM SAGE quantization](trtllm.md#sage-quantization).
@@ -0,0 +1,32 @@
# Hugging Face Hub Attention Backends
Hub backends load the same published kernels commonly used with Hugging Face
Diffusers. They are useful when training and serving must use matching kernel
implementations to reduce numerical drift and sampling divergence.
| Backend | Kernel | Platform |
| --- | --- | --- |
| `FLASH_ATTN_HUB` | `kernels-community/flash-attn2` | Compatible CUDA GPUs |
| `FLASH_ATTN_3_HUB` | `kernels-community/flash-attn3` | Hopper sm_90 or newer |
## Installation
```bash
pip install kernels==0.14.1
```
If `kernels` is unavailable, vLLM-Omni warns and falls back to the
corresponding local FlashAttention implementation. On CUDA devices older than
Hopper, `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
## Usage
```bash
export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB
# Equivalent CLI selection
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN_HUB
```
For per-role kernel alignment, configure the Hub backend on only the relevant
role; see the [attention backend overview](../attention_backends.md#per-role-configuration).
@@ -0,0 +1,72 @@
# RainFusion Attention
`RAINFUSION_ATTN` runs MindIE-SD RainFusion (`rf_v2`) block-sparse video
attention on Ascend NPU. It pools 128-token key blocks, ranks them per query
block, and attends to the highest-scoring blocks after arranging video tokens
in `(t, h, w)` order.
Only the video segment is sparse. Prefix rows such as text, visual conditions,
and audio, plus first-frame blocks, remain dense. Unsupported calls—including
warmup steps, skipped layers, missing video geometry, and video segments below
32 blocks—delegate to `FLASH_ATTN`, so compatible models can select RainFusion
globally.
## Configuration
| Key | Valid values | Meaning |
| --- | --- | --- |
| `sparsity` | finite, `[0, 1]` | Nominal dropped-key-block fraction; default `0.8`; `0` disables sparsity |
| `start_step` | integer, `>= 0` | Number of early denoise steps kept dense |
| `skip_layers` | selector such as `"0-3,38"` | DiT blocks kept dense |
```bash
vllm-omni serve MiniMaxAI/MiniMax-H3 \
--diffusion-attention-config '{"default":{"backend":"RAINFUSION_ATTN",\
"block_sparse":{"sparsity":0.8,"start_step":0}}}'
```
```python
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec
config = AttentionConfig(
default=AttentionSpec(
backend="RAINFUSION_ATTN",
block_sparse=BlockSparseSpec(
sparsity=0.8,
start_step=0,
skip_layers="0-1",
),
),
)
```
Tune `start_step`, then `sparsity`, then `skip_layers`. Increasing the dense
early-step window is usually the cheapest way to recover global structure;
use layer exclusions only after a same-seed dense comparison identifies
sensitive blocks.
## Requirements and compatibility
RainFusion requires Ascend NPU and `mindiesd`. Selecting it on another
platform raises. It is incompatible with ring sequence parallelism because
the kernel needs the complete key sequence for ranking; use Ulysses sequence
parallelism with `ring_degree=1`.
## Geometry handling
RainFusion handles arbitrary video grids by rearranging video spatially and,
when necessary, promoting a real-video suffix to the always-kept prefix. It
does not pad because `rf_v2` does not consume a padding attention mask.
For MiniMax-H3 at 1344x768, grid `(62, 24, 42)` produces 62,496 video rows.
The implementation promotes 2,976 rows and leaves 59,520 sparse rows
(`465 x 128`) so the mask and kernel tiling remain aligned.
Spatial grids aligned to 8x8 are still preferable. A latent height or width
not divisible by 8 creates an always-kept suffix and reduces realized
sparsity. At the image level, that means width and height divisible by 256.
Protected video tails require a compatible MindIE-SD release.
For common configuration and selector behavior, see the
[attention backend overview](../attention_backends.md) and the
[backend selection design](../../../design/feature/attention_backend_selection.md).
@@ -0,0 +1,63 @@
# SageAttention
SageAttention backends provide lossy low-precision attention for diffusion
models. Validate output quality against `TORCH_SDPA` at the same seed before
using either backend in production.
## `SAGE_ATTN`
`SAGE_ATTN` uses SageAttention 2.2 with INT8-quantized attention and FP16
accumulation.
### Installation
Install SageAttention into the same environment as vLLM-Omni:
```bash
git clone https://github.com/thu-ml/SageAttention.git
cd SageAttention
export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
pip install . --no-build-isolation
```
Verify the installation:
```bash
python -c "import sageattention; print(sageattention.__file__)"
```
Select it globally:
```bash
vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN
```
## `SAGE_ATTN_3`
`SAGE_ATTN_3` uses the SageAttention3 Blackwell implementation.
### SageAttention3 installation
```bash
git clone https://github.com/thu-ml/SageAttention.git
cd SageAttention/sageattention3_blackwell
python setup.py install
```
Verify the installation:
```bash
python -c "import sageattn3; print(sageattn3.__file__)"
```
```bash
vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN_3
```
`SAGE_ATTN_3` requires CUDA, an importable `sageattn3`, and a Blackwell-class
GPU. Its kernel assumes the query-head count equals the key/value-head count.
GQA and MQA diffusion calls therefore fall back to PyTorch SDPA for
correctness.
For common configuration and platform routing, see the
[attention backend overview](../attention_backends.md).
@@ -0,0 +1,103 @@
# TRTLLM Attention
`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels. It supports dense
BF16 attention plus optional Skip-Softmax sparsity or SAGE quantization.
For common selection and per-role configuration, see the
[attention backend overview](../attention_backends.md).
## Requirements
Selecting `TRTLLM_ATTN` requires:
- datacenter Blackwell, sm_100 or sm_103;
- `head_dim=128`;
- FlashInfer; and
- a model-declared packed or mask-free compatible path.
It is not supported on workstation Blackwell (sm_120/sm_121). An incompatible
explicit selection raises instead of silently changing the backend.
## Dense mode
Select `TRTLLM_ATTN` without `skip_softmax` or `quant` to run dense BF16:
```bash
vllm-omni serve <model> --diffusion-attention-backend TRTLLM_ATTN
```
## Skip-Softmax
Skip-Softmax trades some fidelity for speed by skipping low-contribution
attention work. See the [feature design](../../../design/feature/skip_softmax.md)
for the algorithm and implementation contract.
| Key | Valid values | Meaning |
| --- | --- | --- |
| `target_sparsity` | finite, `[0, 1]` | Point on a model-calibrated curve |
| `threshold` | finite, `>= 0` | Direct threshold; mutually exclusive with `target_sparsity` |
| `disabled_until_timestep` | finite, `[0, 1]` | Keep early high-noise steps dense until normalized `t <= D` |
```bash
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
"skip_softmax":{"target_sparsity":0.65,"disabled_until_timestep":0.86}}}'
```
```python
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
config = AttentionConfig(
default=AttentionSpec(
backend="TRTLLM_ATTN",
skip_softmax=SkipSoftmaxSpec(
target_sparsity=0.65,
disabled_until_timestep=0.86,
),
),
)
```
Start with `target_sparsity=0.65` and `disabled_until_timestep=0.86`. Increase
sparsity only after comparing quality against dense output at the same seed.
The benefit grows with sequence length because only attention is accelerated.
## SAGE quantization
TRTLLM SAGE quantizes Q/K per block and V per channel so both attention
matrix multiplications use low precision.
| Key | Valid values | Meaning |
| --- | --- | --- |
| `dtype_qk` | `int8`, `fp8_e4m3` | Q/K quantization dtype; absent means dense |
| `q_block_size` | `1`, `4`, `16` | Q scale block size; default `1` |
| `k_block_size` | `1`, `4`, `16` | K scale block size; default `16` |
V is always quantized per channel to FP8 E4M3 and K-smoothing is internal.
Only block sizes with compiled kernels are accepted.
```bash
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
"quant":{"dtype_qk":"fp8_e4m3","q_block_size":1,"k_block_size":16}}}'
```
```python
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
config = AttentionConfig(
default=AttentionSpec(
backend="TRTLLM_ATTN",
quant=AttnQuantSpec(
dtype_qk="fp8_e4m3",
q_block_size=1,
k_block_size=16,
),
),
)
```
This path requires FlashInfer 0.6.16rc1 or newer. FP8 Q/K kernels are
available on sm_100 and sm_103; INT8 Q/K kernels are compiled for sm_100 only.
The shared `AttnQuantSpec` is backend typed: values intended for
`FLASHINFER_ATTN` are rejected here and vice versa.
+60 -374
View File
@@ -1,389 +1,75 @@
# CPU Offloading for Diffusion Models
## Overview
vLLM-Omni provides three CPU-offload strategies for diffusion models. Choose
the coarsest strategy that meets your memory target: finer-grained streaming
saves more device memory but adds more transfer and setup overhead.
vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models:
For the shared factory, discovery, and lifecycle contract, see the
[CPU Offloading design](../../design/feature/offloader/README.md).
1. **Model-level (Sequential) Offloading**: Mutual exclusion between DiT model and encoder - only one is on GPU at a time.
2. **Layerwise (Blockwise) Offloading**: Keeps only one transformer block on GPU at a time with compute-memory overlap.
## Choose a strategy
Both strategies use pinned memory for faster CPU-GPU transfers. The strategies are **mutually exclusive** for now - if both are enabled, layerwise takes priority.
| Strategy | Device residency | Parallel scope | Guide |
| --- | --- | --- | --- |
| Model-level (sequential) | One pipeline component group at a time | Single device | [Model-Level Offloading](offloader/cpu_offload.md) |
| Layerwise (blockwise) | One transformer block, with next-block prefetch | Single device | [Layerwise Offloading](offloader/layerwise_offload.md) |
| Distributed layerwise | Fixed two-block device buffer; optional host-weight sharding and AllGather | Multiple GPU/NPU ranks | [Distributed Layerwise Offloading](offloader/distributed_layerwise_offload.md) |
All strategies use pinned host memory for faster transfers where applicable.
Configuration priority is:
1. Distributed layerwise offloading.
2. Layerwise offloading.
3. Model-level offloading.
Treat the flags as mutually exclusive. If more than one is enabled, the
higher-priority strategy is selected.
## Quick selection
- Use [model-level offloading](offloader/cpu_offload.md) when swapping whole
encoders and DiTs is enough to fit and phase-boundary transfers are
acceptable.
- Use [layerwise offloading](offloader/layerwise_offload.md) for compute-heavy video
DiTs where block transfers can overlap computation.
- Use [distributed layerwise offloading](offloader/distributed_layerwise_offload.md)
when a multi-rank deployment also needs bounded device residency and,
optionally, sharded host weights.
## Supported models
| Architecture | Example models | DiT class | Model-level | Layerwise | Distributed layerwise | Layerwise block attributes |
| --- | --- | --- | :---: | :---: | :---: | --- |
| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | yes | yes | — | `transformer_blocks`, `single_transformer_blocks` |
| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | — | yes | — | `transformer_blocks`, `single_transformer_blocks` |
| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | — | yes | — | `layers` |
| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | — | yes | — | `transformer` |
| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | yes | yes | — | `transformer_blocks` |
| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | yes | yes | — | `down_blocks`, `up_blocks` |
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | — | yes | — | `transformer_blocks` |
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` | yes | yes | — | `layers` |
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | — | yes | — | `layers`, customized modules |
| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | yes | yes | yes | `layers`, `gen_layers` |
Model-level support requires discoverable DiT and encoder components.
Layerwise support requires transformer block topology. Distributed support
reuses that topology but still requires validation for the model, checkpoint,
and parallel configuration.
## Compatibility anchors
These headings preserve links to sections that moved into dedicated guides.
## Model-level (Sequential) Offloading
### How It Works
Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks:
- **When encoders run**: DiT transformer is offloaded to CPU
- **When DiT runs**: Encoders are offloaded to CPU, if more than one dit models, only one loaded on GPU, others get offloaded to CPU.
- **VAE**: Stays resident on GPU
Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed.
### Usage
**Python API:**
```python
from vllm_omni import Omni
m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True)
```
**CLI:**
```bash
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-cpu-offload
```
### To Support a Model
Implement the `SupportsComponentDiscovery` protocol to declare which
submodules serve as pipeline components (used by offloading, HSDP
sharding, and other framework features):
```python
from typing import ClassVar
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
class MyPipeline(nn.Module, SupportsComponentDiscovery):
_dit_modules: ClassVar[list[str]] = ["transformer"]
_encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
_vae_modules: ClassVar[list[str]] = ["vae"]
_resident_modules: ClassVar[list[str]] = [] # optional
def __init__(self):
super().__init__()
self.transformer = ... # DiT — stays on GPU during denoising
self.text_encoder = ... # Encoder — offloaded to CPU during denoising
self.vision_model = ... # Encoder — offloaded to CPU during denoising
self.vae = ... # VAE — always on GPU
```
- `_dit_modules`: attribute names of denoising submodules (kept on GPU
during the diffusion loop).
- `_encoder_modules`: attribute names of encoder/vision submodules
(offloaded to CPU during the diffusion loop).
- `_vae_modules`: attribute names of VAE(s) (always kept on GPU, not
part of the mutual exclusion hooks).
- `_resident_modules`: attribute names of small submodules that must
stay on GPU during layerwise offloading (e.g. embedders, connectors).
Optional — defaults to `[]`.
All attribute names support dotted paths for nested submodules
(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).
Both DiT and encoder lists are needed because the offload hooks use
mutual exclusion: when one group runs, the other moves to CPU.
### Limitations
- Cold start latency increases
- Adds overhead from CPU-GPU transfers between encoder and denoising phases
- Support single GPU only for now
### Component offloading for split models (e.g. Cosmos3)
Some models split their transformer into mutually-exclusive *components* that run
in different phases of a single forward pass rather than as separate pipeline
components -- e.g. Cosmos3's understanding (reasoner) component runs once per
generation while the generation (generator) component runs every denoising step.
Such models have no separate text encoder to swap against, so the transformer
owns a small model-local offload path and wraps each phase with
`with self._offload_context(name):`
```python
class Cosmos3VFMTransformer(nn.Module):
def forward(self, ...):
with self._offload_context("reasoner"):
... # understanding pass, runs once
with self._offload_context("generator"):
... # denoising pass, runs every step
```
Model-level offloading then keeps exactly one component GPU-resident at a time
(the other on CPU), reusing the same `SequentialOffloadHook` `.to()` movers. The
pipeline opts in by exposing `enable_omni_model_cpu_offload` (which drives the
transformer's `enable_model_cpu_offload` and pins the VAE). Layerwise offloading
works for these models too -- each component declares its own block container via
`_layerwise_offload_blocks_attrs`.
See [Model-Level Offloading](offloader/cpu_offload.md).
## Layerwise (Blockwise) Offloading
### How It Works
Layerwise offloading keeps only one transformer block on GPU at a time.
As each block completes, the next block is prefetched to GPU while the current block is freed. The pre and forward hooks utilized by layerwise offloading apply a separate CUDA stream (`copy_stream`) to overlap weight transfer with computation, and retain flattened tensors in pinned CPU memory for block parameters re-materialization. Encoders, VAE, and non-block DiT modules (embeddings, norms) always stay on GPU.
**Execution Flow:**
| Block | Pre-forward Hook | Forward | Post-forward Hook |
|-------|------------------|---------|-------------------|
| block-0 | Prefetch block-1 (async) | Compute block-0 | Free block-0 |
| block-1 | Prefetch block-2 (async) | Compute block-1 | Free block-1 |
| ... | ... | ... | ... |
| block-(n-1) | **Prefetch block-0** (async) | Compute block-(n-1) | Free block-(n-1) |
Each transformer block has a `LayerwiseOffloadHook` that prefetches the next block before forward and frees the current block after forward.
Layerwise offloading is primarily recommended for large **video generation models** where the compute cost per block is high enough to effectively overlap with memory prefetch operations. For example, Wan2.2 T2V and I2V pipelines.
### Usage
**Python API:**
```python
from vllm_omni import Omni
# Text-to-video
m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_layerwise_offload=True)
# Or image-to-video
m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True)
```
**CLI:**
```bash
# Text-to-video
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-layerwise-offload
# Or image-to-video
vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --enable-layerwise-offload
```
### To Support a Model
Models must define the blocks attribute name for layerwise offloading:
```python
class WanTransformer3DModel(nn.Module):
_layerwise_offload_blocks_attrs = ["blocks"] # Attribute names containing transformer blocks
def __init__(self):
self.blocks = nn.ModuleList([...]) # Transformer blocks
```
For models with multiple block types:
```python
class Flux2Transformer2DModel(nn.Module):
_layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
```
### Limitations
- Cold start latency increases because offloaded components must be moved to CPU
during setup; layerwise offload may add extra weight consolidation and pinning
work.
- Performance depends on compute cost and H2D bandwidth as well
- Support single GPU only for now
See [Layerwise Offloading](offloader/layerwise_offload.md).
## Distributed Layerwise Offloading
### How It Works
Distributed layerwise offloading extends single-GPU layerwise offloading to
multi-device deployments. In the default mode, each rank stores only
**1/dp_size** of the model weights in host memory; full layer weights are
reconstructed at runtime via **AllGather** on a dedicated communication stream,
overlapped with computation via a fixed double-buffer scheme. The
`--dlo-no-use-allgather` mode instead streams standard-loader rank-local
weights independently and does not shard weights across ranks.
For the implementation design and compatibility matrix, see
[Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md).
**Key features:**
- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights,
reconstructed per-layer via `all_gather_into_tensor`
- **Fixed double-buffer**: exactly 2 transformer blocks on each device at any
time, regardless of model size
- **DP multi-concurrency**: N concurrent requests processed in parallel
(AllGather only gathers weight shards, which are request-independent)
- **mmap weight loading**: weights loaded as mmap views pointing to shared OS
page cache, eliminating O(dp_size × model_size) RSS during model creation
- **Platform-agnostic**: works on NVIDIA GPU (CUDA/NCCL) and Ascend NPU
(CANN/HCCL) via vLLM-Omni's platform abstraction
**Execution Flow (per device):**
```
Compute Stream: [Layer N] [Layer N+1] [Layer N+2]
H2D Stream: [H2D Shard N+1] [H2D Shard N+2]
AllGather: [AG N+1] [AG N+2]
Slot Usage: Slot0: Layer N Slot1: Layer N+1
```
### Usage
**CLI:**
```bash
# 4× GPU/NPU with AllGather (recommended)
vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--data-parallel-size 4
# Without DLO AllGather (each rank streams standard-loader rank-local weights)
vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--data-parallel-size 4 \
--dlo-no-use-allgather
# With SP instead of DP (long sequences)
vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--usp 4
```
**Python API:**
```python
from vllm_omni import Omni
m = Omni(
model="/path/to/model",
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True, # default
)
```
### CLI Flags
| Flag | Description | Default |
|------|-------------|---------|
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
| `--data-parallel-size N` | Number of DP ranks; DLO weight-sharding group when AllGather is enabled | `1` |
| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` |
| `--dlo-no-use-allgather` | Use standard-loader rank-local weights; no DLO sharding or AllGather | `false` |
### How model weights are loaded (mmap path)
When DLO + AllGather is active, the offloader:
1. **Saves non-persistent buffers** (e.g. RoPE `inv_freq`, timestep `freqs`)
from the normally-created transformer
2. **Converts to meta device** via `to_empty(device="meta")`, releasing random
initialization weights
3. **Loads checkpoint weights as mmap views** via `safe_open().get_tensor()`,
which return views into the OS page cache (shared across all ranks, 0 RSS)
4. **Calls `post_load_weights()`** to apply model-specific dtype conversions
(e.g. Cosmos3's `time_embedder` → FP32)
5. **Restores non-persistent buffers** from saved copies
This approach requires **zero model-specific code changes** — no pipeline or
transformer modifications are needed.
### OffloadPlan (declarative topology metadata)
Models can optionally declare an `OffloadPlan` class variable to provide
topology metadata (block attribute names, submodules to offload) without
any offload-specific logic:
```python
from vllm_omni.diffusion.offloader import OffloadPlan
class MyPipeline(nn.Module):
_dit_modules = ["transformer"]
_offload_plan = OffloadPlan(
block_attrs={"transformer": ("blocks",)},
offload_submodules={"context_encoder": "layers"},
)
```
When not declared, the offloader falls back to `_layerwise_offload_blocks_attrs`
and heuristic attribute search. This is backward-compatible — existing models
work without any changes.
### DP Multi-concurrency
When `--data-parallel-size > 1` and AllGather is enabled, the scheduler batches
up to `dp_size` requests per denoise step. Each DP rank processes a different
request while AllGather synchronizes only weight shards (request-independent).
**Requirements for concurrent requests:**
- `num_inference_steps` must be specified explicitly (None is not allowed)
- All concurrent requests must have the same `num_inference_steps` value
- These constraints exist because AllGather is a collective that requires
every rank to participate at each denoise step
### Limitations
- Online quantization (FP8) is incompatible with mmap loading — falls back to
regular loading only when the model is outside the DLO+AllGather mmap path;
online quantization with DLO+AllGather is rejected. Use
`--dlo-no-use-allgather` or disable online quantization.
- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware
weight-loader callbacks are bypassed. The no-AllGather path retains the
standard TP loader and is experimental.
- HSDP + AllGather is rejected because it would double-shard parameters. HSDP
with no-AllGather is accepted at configuration level, but has limited
end-to-end validation.
- `num_inference_steps=None` is not allowed in DP multi-concurrency mode
**Module Discovery**
The offloader discovers pipeline components in two ways:
1. **Protocol-based** (preferred): If the pipeline implements
`SupportsComponentDiscovery`, its `_dit_modules`, `_encoder_modules`,
`_vae_modules`, and `_resident_modules` class variables are used
directly. All attribute names support dotted paths (e.g.
`"pipe.transformer"`, `"bagel.time_embedder"`) for nested submodules.
2. **Fallback attribute scan**: Otherwise, the offloader scans for
well-known attribute names:
- **DiT modules**: `transformer`, `transformer_2`, `dit`, `sr_dit`, `language_model`, `transformer_blocks`, `model`
- **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder`
- **VAE**: `vae`, `audio_vae`
**Hook System**
Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelHook`) to register pre/post forward callbacks on modules, enabling automatic swapping without modifying model code.
**Backend Architecture**
```
OffloadBackend (base class)
├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap)
│ (delegates to a pipeline's enable_omni_model_cpu_offload
│ for split models like Cosmos3)
├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
│ (single-GPU, full weights on host)
└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
(multi-GPU, sharded AllGather or rank-local mode)
```
Factory function `get_offload_backend()` selects the appropriate backend based on
configuration.
For split models, `ModelLevelOffloadBackend.enable()` detects a pipeline's
`enable_omni_model_cpu_offload` hook and delegates to it; Cosmos3 then swaps its
reasoner/generator components inside the model forward pass.
## Supported Models
| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Distributed Layerwise Offload | Blocks Attrs (Layerwise specific) |
|--------------|----------------|-----------|---------------------|-------------------|-------------------------------|-----------------------------------|
| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | - | `"layers"` |
| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | - | `"transformer"` |
| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"` |
| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | - | `"down_blocks"`, `"up_blocks"` |
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | - | `"transformer_blocks"` |
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | - | `"layers"` |
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | - | `"layers"`, `"customized modules"` |
| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | ✓ | ✓ | ✓ | `"layers"`, `"gen_layers"` |
**Notes:**
- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
- Distributed Layerwise Offloading reuses the layerwise block topology, but the
model and parallelism combination still needs validation. AllGather mode
additionally requires a compatible checkpoint/mmap layout; no-AllGather mode
retains the standard loader's rank-local tensors. See the
[DLO design](../../design/feature/distributed_layerwise_offload.md) and
[Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
for usage examples.
See [Distributed Layerwise Offloading](offloader/distributed_layerwise_offload.md).
@@ -0,0 +1,81 @@
# Model-Level Offloading
Model-level, or sequential, offloading keeps only the pipeline component group
currently executing on the accelerator. It is the simplest offload strategy
and is selected with `--enable-cpu-offload`.
## How it works
Pre-forward hooks enforce mutual exclusion between DiT and encoder modules:
- before an encoder runs, the DiT moves to CPU;
- before a DiT runs, encoders and other DiTs move to CPU; and
- VAE modules remain on the accelerator.
Pinned host memory reduces transfer overhead. Transfers occur at phase
boundaries, so cold-start and encoder-to-denoiser transitions become slower.
## Usage
```python
from vllm_omni import Omni
omni = Omni(
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
enable_cpu_offload=True,
)
```
```bash
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--omni --enable-cpu-offload
```
## Model integration
Pipelines should implement `SupportsComponentDiscovery`:
```python
from typing import ClassVar
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
class MyPipeline(nn.Module, SupportsComponentDiscovery):
_dit_modules: ClassVar[list[str]] = ["transformer"]
_encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
_vae_modules: ClassVar[list[str]] = ["vae"]
_resident_modules: ClassVar[list[str]] = []
```
All entries may be dotted paths. DiT and encoder lists are both required for
mutual exclusion. VAE modules are pinned but not swapped; resident modules are
small modules that must stay on the accelerator for layerwise paths.
## Split-model components
Some models, such as Cosmos3, split one transformer into mutually exclusive
components that run in different phases. The pipeline exposes
`enable_omni_model_cpu_offload`, and the backend delegates to the model-local
contexts:
```python
class Cosmos3VFMTransformer(nn.Module):
def forward(self, ...):
with self._offload_context("reasoner"):
...
with self._offload_context("generator"):
...
```
This preserves the same invariant—exactly one component is device resident—
while reusing sequential `.to()` movers.
## Limitations
- Single device only.
- Higher cold-start latency.
- Transfers between encoder and denoising phases add latency.
See the [model-level design](../../../design/feature/offloader/cpu_offload.md)
for lifecycle and extension invariants.
@@ -0,0 +1,118 @@
# Distributed Layerwise Offloading
Distributed layerwise offloading (DLO) extends block streaming to multi-device
deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
the host weights and reconstructs each layer at runtime. Without AllGather,
each rank streams its standard-loader rank-local weights independently.
See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
for the implementation contract and compatibility matrix.
## Execution model
DLO overlaps three operations with a fixed two-block device buffer:
```text
Compute stream: [Layer N] [Layer N+1] [Layer N+2]
H2D stream: [H2D shard N+1] [H2D shard N+2]
AllGather: [AG N+1] [AG N+2]
Slots: slot 0: Layer N slot 1: Layer N+1
```
AllGather communicates only request-independent weight shards, so data-
parallel ranks may process different requests concurrently.
## Usage
```bash
# Four ranks with sharded host weights and AllGather
vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--data-parallel-size 4
# Standard-loader rank-local weights, without DLO AllGather
vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--data-parallel-size 4 \
--dlo-no-use-allgather
# Sequence parallel deployment
vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--usp 4
```
```python
from vllm_omni import Omni
omni = Omni(
model="/path/to/model",
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
)
```
## Flags
| Flag | Meaning | Default |
| --- | --- | --- |
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
| `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
| `--dlo-no-use-allgather` | Keep standard-loader rank-local weights | `false` |
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
## mmap weight loading
The DLO plus AllGather path:
1. saves non-persistent buffers such as RoPE frequencies;
2. moves the normally created transformer to the meta device;
3. loads checkpoint tensors as mmap views backed by the shared OS page cache;
4. calls model-specific `post_load_weights()` conversions; and
5. restores the saved non-persistent buffers.
This avoids one full checkpoint RSS copy per rank and does not require
model-specific loading code.
## Declarative topology
Models may declare an `OffloadPlan` instead of embedding offload logic:
```python
from vllm_omni.diffusion.offloader import OffloadPlan
class MyPipeline(nn.Module):
_dit_modules = ["transformer"]
_offload_plan = OffloadPlan(
block_attrs={"transformer": ("blocks",)},
offload_submodules={"context_encoder": "layers"},
)
```
When no plan exists, discovery falls back to
`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup.
## Data-parallel concurrency
With `data_parallel_size > 1` and AllGather enabled, the scheduler can process
up to `dp_size` requests per denoising step. Every concurrent request must set
the same explicit `num_inference_steps`; `None` is rejected because every rank
must enter each collective.
## Limitations
- Online FP8 quantization is rejected with the DLO plus AllGather mmap path.
Use `--dlo-no-use-allgather` or disable online quantization.
- Tensor parallel size greater than one is rejected in the mmap path because
it bypasses TP-aware loader callbacks. The no-AllGather path retains the
standard loader and remains experimental with TP.
- HSDP plus AllGather is rejected to avoid double sharding. HSDP without
AllGather has limited end-to-end validation.
- Resident leading layers require `--dlo-no-use-allgather` and a model
`OffloadPlan` that declares eligible `resident_dit_paths`.
- DP concurrency requires an explicit, identical inference-step count.
See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
for an end-to-end example.
@@ -0,0 +1,65 @@
# Layerwise Offloading
Layerwise, or blockwise, offloading keeps one transformer block on the
accelerator and prefetches the next block while the current block computes.
It is best suited to compute-heavy video DiTs whose block execution can hide
host-to-device transfers.
## Execution flow
Each block has a pre-forward and post-forward hook. Parameters are consolidated
in pinned host tensors and rematerialized for execution on a dedicated copy
stream.
| Block | Pre-forward hook | Forward | Post-forward hook |
| --- | --- | --- | --- |
| block 0 | Prefetch block 1 | Compute block 0 | Free block 0 |
| block 1 | Prefetch block 2 | Compute block 1 | Free block 1 |
| ... | ... | ... | ... |
| last block | Prefetch block 0 | Compute last block | Free last block |
Encoders, VAE modules, and non-block DiT modules such as embeddings and norms
remain device resident.
## Usage
```python
from vllm_omni import Omni
omni = Omni(
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
enable_layerwise_offload=True,
)
```
```bash
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--omni --enable-layerwise-offload
```
## Model integration
Transformer classes declare containers of executable blocks:
```python
class WanTransformer3DModel(nn.Module):
_layerwise_offload_blocks_attrs = ["blocks"]
class Flux2Transformer2DModel(nn.Module):
_layerwise_offload_blocks_attrs = [
"transformer_blocks",
"single_transformer_blocks",
]
```
See the [layerwise design](../../../design/feature/offloader/layerwise_offload.md) for
the discovery and hook invariants. `OffloadPlan` is a separate declarative
topology path for distributed layerwise offload.
## Limitations
- Single device only.
- Setup consolidates and pins block parameters, increasing cold-start time.
- Performance depends on block compute time and host-to-device bandwidth;
lightweight blocks may not hide transfers.
@@ -64,6 +64,41 @@ omni = Omni(
)
```
### Combined with Online FP8 Quantization
HSDP can be combined with online FP8 quantization for diffusion models that
support both features. HSDP shards the quantized weights across GPUs to reduce
the per-GPU weight memory footprint.
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="<your-model>",
quantization="fp8",
parallel_config=DiffusionParallelConfig(
use_hsdp=True,
hsdp_shard_size=2,
),
)
```
For online serving:
```bash
vllm serve <your-model> \
--omni \
--quantization fp8 \
--use-hsdp \
--hsdp-shard-size 2
```
The model must support both HSDP and online FP8 quantization. HSDP remains
incompatible with Tensor Parallelism. See the
[FP8 Quantization Guide](../../quantization/fp8.md) for supported models and
hardware.
---
## Example Script
@@ -147,3 +182,4 @@ For detailed instructions on adding HSDP support to new models, see the [HSDP Co
1. ✅ **Enable HSDP** - Set `use_hsdp=True` and `hsdp_shard_size` to reduce per-GPU memory for large models
2. ✅ **Combine with SP** - Use together with `ulysses_degree` for video models requiring both memory reduction and sequence parallelism
3. ⚠️ **Incompatible with TP** - `tensor_parallel_size` must be 1 when HSDP is enabled
4. ✅ **Combine with online FP8** - Use `quantization="fp8"` or `--quantization fp8` to reduce the sharded weight memory footprint
+5 -2
View File
@@ -106,7 +106,10 @@ The following tables show which models support each feature:
> Notes:
> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload.md).
> 1. CPU Offload has three strategies: model-level, layerwise, and distributed
> layerwise. The tables below show **layerwise support** only. Split models
> like Cosmos3 swap their reasoner/generator components for model-level
> offload; see the [CPU Offload Guide](diffusion/cpu_offload.md).
> 2. The **💾Quantization** column is collapsed for readability. See [Quantization](quantization/overview.md) for per-method and per-model support details.
### ImageGen
@@ -214,7 +217,7 @@ The following tables show which models support each feature:
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).
compatibility matrix in the [Distributed Layerwise Offloading guide](diffusion/offloader/distributed_layerwise_offload.md).
5. The compatibility matrix uses FP8 as the representative quantization method.
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).
+1 -1
View File
@@ -165,7 +165,7 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \
- **Tensor Parallelism — Text Encoder Not Sharded**: TP currently only shards the DiT blocks. Each TP rank retains a **full copy of the text encoder weights**, leading to significant GPU memory overhead proportional to TP degree. Tracked in [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).
- **CPU Offloading — Two Modes Are Mutually Exclusive**: Model-level offload (`enable_cpu_offload`) and layerwise offload (`enable_layerwise_offload`) cannot be used simultaneously. If both are set, layerwise takes priority and model-level is silently ignored.
- **CPU Offloading — Strategies Are Mutually Exclusive**: Distributed layerwise, layerwise, and model-level offload should not be enabled together. Selection priority is distributed layerwise, then layerwise, then model-level.
- **CPU Offloading — VAE stays on GPU**: Both offloading strategies keep the VAE on GPU at all times. For high-resolution generation, VAE decode can still cause OOM. Mitigate by combining with `vae_use_tiling=True` or VAE Patch Parallelism.
+3 -1
View File
@@ -80,12 +80,14 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
## Model Type Support
### Diffusion Model (Qwen-Image, Wan2.2)
### Diffusion Models
| Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization |
|-------|-----------|:-------:|:------:|----------------|------------------|------------------|
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | |
| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | |
| LTX-2 | `Lightricks/LTX-2`, `rootonchair/LTX-2-19b-distilled` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers`, `diffusers/LTX-2.3-Distilled-Diffusers` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None | ✅︎ |
| FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None | |
| FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | Yes | Not validated | All layers | None | ✅︎ |
@@ -59,7 +59,7 @@ def main(args):
omni = Omni(
model=args.model_name,
stage_configs_path=args.stage_configs_path,
deploy_config=args.deploy_config,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
batch_timeout=args.batch_timeout,
@@ -167,10 +167,10 @@ def parse_args():
help="Sampling rate for audio loading (default: 16000).",
)
parser.add_argument(
"--stage-configs-path",
"--deploy-config",
type=str,
default=None,
help="Path to stage configs YAML file.",
help="Path to a deploy config YAML file.",
)
parser.add_argument(
"--log-stats",
@@ -876,7 +876,7 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
parser.add_argument("--task", type=str, required=True, choices=TASK_CHOICES)
parser.add_argument("--model", type=str, required=True, help="HF repo id or local model directory.")
parser.add_argument(
"--stage-config-path",
"--deploy-config",
type=str,
default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
help="Path to deploy config yaml.",
@@ -1394,8 +1394,8 @@ def main() -> None:
from vllm_omni.entrypoints.omni import Omni
stage_config_path = str(Path(args.stage_config_path).expanduser())
omni = Omni(model=model_source, stage_configs_path=stage_config_path, dtype=args.dtype)
deploy_config_path = str(Path(args.deploy_config).expanduser())
omni = Omni(model=model_source, deploy_config=deploy_config_path, dtype=args.dtype)
sampling_params_list = [
SamplingParams(max_tokens=int(args.max_tokens_per_stage), temperature=0.0, top_p=1.0, detokenize=False)
for _ in range(omni.num_stages)
@@ -73,7 +73,6 @@ def parse_args():
)
parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.")
parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.")
parser.add_argument(
"--stream",
action="store_true",
@@ -146,13 +145,7 @@ def main():
else:
bot_task = args.bot_task
if args.deploy_config is not None and args.stage_configs_path is not None:
raise ValueError("--deploy-config and --stage-configs-path are mutually exclusive.")
deploy_config = args.deploy_config
stage_configs_path = args.stage_configs_path
if deploy_config is None and stage_configs_path is None:
deploy_config = _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
deploy_config = args.deploy_config or _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
omni_kwargs = {
"model": args.model,
@@ -168,10 +161,7 @@ def main():
if additional_config is not None:
omni_kwargs["additional_config"] = additional_config
if deploy_config is not None:
omni_kwargs["deploy_config"] = deploy_config
else:
omni_kwargs["stage_configs_path"] = stage_configs_path
omni_kwargs["deploy_config"] = deploy_config
omni = Omni(**omni_kwargs)
@@ -272,10 +262,7 @@ def main():
print(f" Modality: {args.modality}")
print(f" Prompt task: {task}")
print(f" Bot task: {bot_task}")
if deploy_config is not None:
print(f" Deploy config: {deploy_config}")
else:
print(f" Stage config: {stage_configs_path}")
print(f" Deploy config: {deploy_config}")
print(f" Num stages: {omni.num_stages}")
if args.modality in ("text2img", "img2img"):
print(f" Inference steps: {args.steps}")
@@ -452,12 +452,6 @@ def parse_args():
default=None,
help="Path to a .txt file with one prompt per line (preferred).",
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="Path to a stage configs file (deprecated; prefer --deploy-config).",
)
parser.add_argument(
"--deploy-config",
type=str,
@@ -484,10 +484,10 @@ def parse_args():
help="Path to a .txt file with one prompt per line (preferred).",
)
parser.add_argument(
"--stage-configs-path",
"--deploy-config",
type=str,
default=None,
help="Path to a stage configs file.",
help="Path to a deploy config YAML.",
)
parser.add_argument(
"--video-path",
@@ -232,15 +232,8 @@ def main(args):
raise ValueError(f"Unknown query type: {args.query_type}")
# Initialize vLLM-Omni with Step-Audio2
# Resolve deploy config path. Keep explicit stage_configs_path for legacy
# custom configs, but use bundled deploy configs by default.
deploy_config_path = None
stage_config_path = None
if args.stage_configs_path and args.deploy_config:
raise ValueError("--stage-configs-path and --deploy-config are mutually exclusive")
if args.stage_configs_path:
stage_config_path = args.stage_configs_path
elif args.deploy_config:
# Resolve an explicit deploy config or use the bundled task default.
if args.deploy_config:
deploy_config_path = args.deploy_config
else:
configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/deploy"
@@ -262,10 +255,7 @@ def main(args):
"ray_address": args.ray_address,
"trust_remote_code": True,
}
if stage_config_path is not None:
omni_kwargs["stage_configs_path"] = stage_config_path
else:
omni_kwargs["deploy_config"] = deploy_config_path
omni_kwargs["deploy_config"] = deploy_config_path
omni_llm = Omni(**omni_kwargs)
@@ -418,12 +408,6 @@ def parse_args():
" - HuggingFace ID: stepfun-ai/Step-Audio-2-mini"
),
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="Path to a legacy stage config YAML file (default: use bundled deploy config)",
)
parser.add_argument(
"--deploy-config",
type=str,
@@ -92,12 +92,6 @@ def parse_args() -> argparse.Namespace:
"meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
"stabilityai/stable-diffusion-3.5-medium, Tongyi-MAI/Z-Image-Turbo and etc.",
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="[Deprecated] Path to a legacy stage_args-format YAML. Prefer --deploy-config.",
)
parser.add_argument(
"--deploy-config",
type=str,
@@ -475,8 +469,6 @@ def main():
omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
if args.enforce_eager is not None:
omni_kwargs["enforce_eager"] = args.enforce_eager
if args.stage_configs_path:
omni_kwargs["stage_configs_path"] = args.stage_configs_path
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
if use_nextstep:
@@ -515,8 +507,8 @@ def main():
print(f" Image size: {args.width}x{args.height}")
if args.lora_path:
print(f" LoRA: scale={args.lora_scale}")
if args.stage_configs_path:
print(f" stage-configs-path: {args.stage_configs_path}")
if args.deploy_config:
print(f" deploy-config: {args.deploy_config}")
print(f"{'=' * 60}\n")
# Build LoRA request when --lora-path is set
@@ -23,6 +23,7 @@ list of supported architectures across all modalities, see
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz |
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz |
| dots.tts | `rednote-hilab/dots.tts-soar` | single (native AR) | — (not wired yet) | — | — | 48 kHz |
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz |
| IndexTTS-2.5 | native `checkpoints/` bundle | 2 (AR talker + EnhancedCodec + S2Mel DiT + BigVGAN) | ✓ (required) | — | multilingual (`--lang`) + emotion control | 22.05 kHz |
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
@@ -430,6 +431,31 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [
---
## dots.tts
Single-stage native AR TTS at 48 kHz (rednote-hilab). Pipeline: `Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR loopback → AudioVAE (streaming decode)`. Same "vLLM-native base LM + side-path computation" pattern as VoxCPM2, with a plain Qwen2 backbone instead of MiniCPM4 and no FSQ / residual LM stage.
### Quick start
```bash
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
--model rednote-hilab/dots.tts-soar \
--text "Hello, this is a test of dots TTS running on vLLM Omni."
```
### Voice cloning
Not wired in this release — generation is zero-shot only. The CAM++ x-vector speaker encoder weights load, but `end2end.py` has no `--ref-audio`/`--ref-text` flags yet.
### Streaming
The AudioVAE decoder has an internal streaming path (`init_stream_state` / `stream_step` / `stream_flush`) used to avoid boundary artifacts between 160 ms patches, but it is not yet exposed through an online serving endpoint or example.
### Notes
- Output: 48 kHz mono WAV.
- Deploy config: `vllm_omni/deploy/dots_tts.yaml` (auto-loaded by HF `model_type`).
- Checkpoints: `rednote-hilab/dots.tts-soar` is the validated default. `dots.tts-base` shares the same architecture but is unvalidated in this repo. `dots.tts-mf` (MeanFlow, 2-4 step) is not supported yet.
- Known limitation: no CUDA graph capture and no batched side-path computation yet, so concurrent requests do not currently scale (each request's DiT Euler steps run serially). See `recipes/rednote-hilab/dots.tts.md` for details and the roadmap.
---
## IndexTTS-2 and IndexTTS-2.5
Both versions use a 2-stage TTS pipeline and produce 22.05 kHz mono audio. Stage 0 is a GPT AR talker. IndexTTS-2 Stage 1 uses RepCodec semantic embeddings, GPT latent, S2Mel CFM/DiT, and BigVGAN. IndexTTS-2.5 replaces RepCodec with EnhancedCodec and uses the official code-only path (`use_gpt_latent=false`). Every request requires reference audio for zero-shot voice cloning.
@@ -0,0 +1,125 @@
"""Offline dots.tts inference example (single-stage native AR pipeline).
Uses the single-stage native AR config (vllm_omni/deploy/dots_tts.yaml).
Reference checkpoint: rednote-hilab/dots.tts-soar.
"""
from __future__ import annotations
import os
import time
from pathlib import Path
import soundfile as sf
import torch
from vllm_omni import Omni
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
REPO_ROOT = Path(__file__).resolve().parents[4]
SAMPLE_RATE = 48_000
def parse_args():
parser = TrackingArgumentParser(description="Offline dots.tts native AR inference")
parser.add_argument(
"--model",
type=str,
default="rednote-hilab/dots.tts-soar",
help="dots.tts model path or HuggingFace repo ID.",
)
parser.add_argument(
"--text",
type=str,
default="Hello, this is a test of dots TTS running on vLLM Omni.",
help="Text to synthesize.",
)
parser.add_argument(
"--output-dir",
type=str,
default="output_audio",
help="Directory for output WAV files.",
)
parser.add_argument(
"--deploy-config",
type=str,
default=None,
help="Override the deploy config path. If unset, auto-loads "
"vllm_omni/deploy/dots_tts.yaml based on the HF model_type.",
)
return parser.parse_args()
def extract_audio(multimodal_output: dict) -> torch.Tensor:
"""Extract the final complete audio tensor from multimodal output.
The output processor concatenates per-step delta tensors under
``model_outputs``. Falls back to ``audio`` for backwards compat.
"""
audio = multimodal_output.get("model_outputs")
if audio is None:
audio = multimodal_output.get("audio")
if audio is None:
raise ValueError(f"No audio key in multimodal_output: {list(multimodal_output.keys())}")
if isinstance(audio, list):
valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
if not valid:
raise ValueError("Audio list is empty or all elements are None.")
return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
return torch.as_tensor(audio).float().cpu().reshape(-1)
def main():
args = parse_args()
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
engine = Omni(
model=args.model,
deploy_config=args.deploy_config,
)
from transformers import AutoTokenizer
from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import (
build_dots_tts_prompt,
)
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
prompt = build_dots_tts_prompt(
tokenizer=tokenizer,
text=args.text,
)
print(f"Model : {args.model}")
print(f"Text : {args.text}")
print(f"Output dir : {output_dir}")
print(f"Prompt len : {len(prompt['prompt_token_ids'])} tokens")
t_start = time.perf_counter()
outputs = engine.generate([prompt])
elapsed = time.perf_counter() - t_start
request_output = outputs[0]
mm = request_output.outputs[0].multimodal_output
audio = extract_audio(mm)
duration = audio.numel() / SAMPLE_RATE
rtf = elapsed / duration if duration > 0 else float("inf")
output_path = output_dir / "output.wav"
sf.write(str(output_path), audio.numpy(), SAMPLE_RATE, format="WAV")
print(f"Saved : {output_path}")
print(f"Duration : {duration:.2f}s")
print(f"Inference : {elapsed:.2f}s")
print(f"RTF : {rtf:.3f}")
if __name__ == "__main__":
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
main()
@@ -38,7 +38,7 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
logger = logging.getLogger(__name__)
DEFAULT_MODEL = "fishaudio/s2-pro"
DEFAULT_STAGE_CONFIG = os.path.join(
DEFAULT_DEPLOY_CONFIG = os.path.join(
os.path.dirname(__file__),
"..",
"..",
@@ -152,10 +152,7 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
else:
omni_kwargs["stage_configs_path"] = args.stage_configs_path or DEFAULT_STAGE_CONFIG
omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG
return omni_kwargs
@@ -278,12 +275,6 @@ def parse_args():
default=None,
help="Deploy config YAML path (under vllm_omni/deploy/).",
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help="Path to stage configs YAML (legacy, prefer --deploy-config).",
)
parser.add_argument(
"--output-dir",
default="output_audio",
@@ -145,7 +145,7 @@ def main(args):
omni = Omni(
model=args.model,
stage_configs_path=deploy_config_path,
deploy_config=deploy_config_path,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
@@ -490,10 +490,10 @@ def parse_args():
help="Path to a .txt file with one prompt per line (preferred).",
)
parser.add_argument(
"--stage-configs-path",
"--deploy-config",
type=str,
default=None,
help="Path to a stage configs file.",
help="Path to a deploy config YAML.",
)
parser.add_argument(
"--audio-path",
+6
View File
@@ -73,12 +73,18 @@ plugins:
- awesome-nav
- redirects:
redirect_maps:
community/meetups.md: community/meetings.md
community/volunteers.md: community/contact_us.md
user_guide/diffusion/request_batching.md: user_guide/diffusion/execution_modes.md
user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
user_guide/quantization/index.md: user_guide/quantization/overview.md
design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
design/feature/cpu_offloading.md: design/feature/offloader/README.md
design/feature/model_level_offload.md: design/feature/offloader/cpu_offload.md
design/feature/layerwise_offload.md: design/feature/offloader/layerwise_offload.md
design/feature/distributed_layerwise_offload.md: design/feature/offloader/distributed_layerwise_offload.md
- glightbox
- git-revision-date-localized:
# exclude files
+16
View File
@@ -40,6 +40,7 @@ dev = [
"pytest-asyncio==1.4.0",
"pytest-cov==7.1.0",
"pytest-mock==3.15.1",
"pytest-shard==0.1.2",
"websockets==16.1",
"pytest-subtests==0.15.0",
"datasets==3.6.0",
@@ -301,6 +302,21 @@ filterwarnings = [
"ignore:.*does not have '__bases__' attribute.*:UserWarning",
]
[tool.coverage.run]
core = "sysmon"
source = ["vllm_omni"]
# Online-serving tests launch the server via subprocess.Popen and stop it with
# SIGTERM. pytest-cov no longer instruments subprocesses implicitly, and a
# SIGTERM-ed process discards its data, so without patch/sigterm a coverage run
# over an online test measures only the pytest parent.
patch = ["subprocess"]
concurrency = ["multiprocessing", "thread"]
sigterm = true
[tool.coverage.report]
show_missing = true
skip_covered = true
[tool.typos.default]
extend-ignore-identifiers-re = [
".*_thw",
+72 -46
View File
@@ -4,11 +4,20 @@
## Pipelines
| Model | Checkpoint | Task | `--model-class-name` | Batching |
|---|---|---|---|---|
| LTX-2 | `Lightricks/LTX-2` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
| LTX-2 distilled | `rootonchair/LTX-2-19b-distilled` | Two-stage T2V/I2V | `LTX2DistilledPipeline` | No |
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
| `--model-class-name` | Task | Required checkpoint repositories |
|---|---|---|
| `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` |
| `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` |
| `LTX2DistilledPipeline` | LTX-2 full-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
| `LTX2Pipeline` | LTX-2.3 one-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` |
| `LTX2TwoStagePipeline` | LTX-2.3 ordinary two-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers`<br>`Lightricks/LTX-2.3` |
| `LTX2DistilledPipeline` | LTX-2.3 full-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`<br>`Lightricks/LTX-2.3` |
Repositories in the table are download units. A full pipeline repository
contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler,
and tokenizer; an additional repository supplies LoRA or upsampler sidecars.
The runtime first searches the model root for sidecars under their official
filenames, then downloads them from the matching Lightricks Hub repository.
`LTX2Pipeline` is the unified one-stage entry. Checkpoint metadata selects the
LTX-2 or LTX-2.3 profile; omitting an image selects T2V, while one initial
@@ -16,10 +25,10 @@ image selects I2V. Both one-stage repositories declare this class, so
`--model-class-name` is optional. LTX-2.3 requires the Diffusers checkpoint;
the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable.
`LTX2DistilledPipeline` is the unified distilled entry: omit `image` for T2V
or provide one initial image for I2V. As in the official LTX pipelines,
distilled mode is selected explicitly rather than inferred from the checkpoint
name or metadata; always pass `--model-class-name LTX2DistilledPipeline`.
`LTX2TwoStagePipeline` samples the regular model at half resolution, upsamples,
then refines with the distilled LoRA. `LTX2DistilledPipeline` uses a fully
merged distilled Transformer in both stages. Both entries support T2V and I2V;
select their class explicitly.
## API Migration
@@ -68,24 +77,28 @@ The model recipe also supplies the default negative prompt. Top-level
API defaults `num_frames` to `1`, so set it explicitly; the offline LTX scripts
default to `121`.
## Distilled Two-Stage Defaults
## Two-Stage Defaults
| Parameter | Value |
|---|---:|
| Final width × height | 1536 × 1024 |
| Stage 1 width × height | 768 × 512 |
| Frames / frame rate | 121 / 24 |
| Stage 1 / Stage 2 denoise steps | 8 / 3 |
| Guidance | Fixed positive-only |
| Parameter | Ordinary | Full-distilled |
|---|---:|---:|
| Final width × height | 1536 × 1024 | 1536 × 1024 |
| Stage 1 width × height | 768 × 512 | 768 × 512 |
| Frames / frame rate | 121 / 24 | 121 / 24 |
| Stage 1 / Stage 2 steps | 40 (LTX-2) or 30 (LTX-2.3) / 3 | 8 / 3 |
| Guidance | Stage 1 guided; Stage 2 positive-only | Fixed positive-only |
The dimensions passed to the API are the final output dimensions and must be
divisible by 64. Both one- and two-stage requests require `num_frames = 8k+1`.
The distilled recipe fixes both sigma schedules; `num_inference_steps` must be
`8` when supplied.
API dimensions are final dimensions and must be divisible by 64. All LTX
requests require `num_frames = 8k+1`. Ordinary Stage 1 uses the LTX-2 or
LTX-2.3 one-stage defaults shown above. Distilled schedules are fixed, so
`num_inference_steps`, when supplied, must be `8`. Both entries reject custom
sigmas and input latents.
Ordinary two-stage uses layer-fused LoRA for an unquantized BF16 Transformer
and automatically switches to dynamic LoRA when quantization is enabled.
## Serving
Start either checkpoint:
Start either one-stage checkpoint:
```bash
vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
@@ -95,11 +108,19 @@ vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
vllm serve diffusers/LTX-2.3-Diffusers --omni --stage-init-timeout 600
```
Start the distilled checkpoint with its two-stage entry:
Start a two-stage checkpoint with its explicit class:
```bash
vllm serve rootonchair/LTX-2-19b-distilled --omni \
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
# LTX-2.3 full-distilled; the v1.1 x2 upsampler is resolved when absent
vllm serve diffusers/LTX-2.3-Distilled-Diffusers --omni \
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
# Ordinary LTX-2.3 two-stage
vllm serve diffusers/LTX-2.3-Diffusers --omni \
--model-class-name LTX2TwoStagePipeline \
--enable-layerwise-offload \
--stage-init-timeout 600
```
The same server handles T2V and I2V. A T2V request using the selected recipe's
@@ -135,8 +156,9 @@ together with `input_reference`.
## Guidance
LTX supports independent video/audio CFG, spatio-temporal guidance (STG),
cross-modality guidance, and rescaling:
One-stage and ordinary Stage 1 support independent video/audio CFG,
spatio-temporal guidance (STG), cross-modality guidance, and rescaling.
Distilled stages and ordinary Stage 2 are fixed positive-only.
| Parameter | Default | Effect | Alias |
|---|---:|---|---|
@@ -172,9 +194,9 @@ complete guidance plan in parallel. Despite the legacy `cfg` name, it covers
text CFG, STG, and cross-modality guidance passes; guidance rescaling is
applied after all pass predictions are gathered.
The default one-stage recipe has four Transformer passes per denoise step:
`cond`, `uncond`, `ptb` (STG), and `mod` (cross-modality). The useful balanced
configurations are therefore:
The default one-stage recipe and ordinary Stage 1 have four Transformer passes
per denoise step: `cond`, `uncond`, `ptb` (STG), and `mod`
(cross-modality). The useful balanced configurations are therefore:
| `--cfg-parallel-size` | Passes per rank | Guidance-slot utilization | Notes |
|---:|---:|---:|---|
@@ -195,9 +217,8 @@ vllm serve diffusers/LTX-2.3-Diffusers --omni \
```
The total device count is the product of `cfg_parallel_size` and the other
configured parallel dimensions. Distilled pipelines use positive-only
guidance, so `cfg_parallel_size > 1` does not provide guidance speedup for
them.
configured parallel dimensions. Positive-only distilled phases do not benefit
from `cfg_parallel_size > 1`.
### Python API
@@ -242,22 +263,22 @@ noted below.
| `image` | image or batch, `None` | Direct value wins over request images; no image selects T2V. I2V accepts one image per prompt, and a batch cannot mix T2V/I2V. |
| `prompt` | string or list, `None` | Positive-text fallback; request prompts win. Mutually exclusive with `prompt_embeds`. |
| `negative_prompt` | string or list, `None` | Fallback after request values and before the recipe default; mutually exclusive with negative embeddings. |
| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 distilled. |
| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 two-stage. |
| `width` | `int`, `None` | Same precedence and alignment as `height`. |
| `num_frames` | `int`, `None` | Request → direct value → recipe default; must be `8k+1` and also determines audio duration with `frame_rate`. |
| `frame_rate` | `float`, `None` | Request `frame_rate` → request `fps` → direct value → recipe default. |
| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. |
| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage and ordinary Stage 1, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. |
| `sigmas` | list of float, `None` | One-stage only. Request values win; every request in a fused batch must use the same schedule. |
| `timesteps` | list of int, `None` | Compatibility slot; LTX accepts only `None`. Use `sigmas`. |
| `guidance_scale` | `float`, `None` | One-stage common video/audio CFG fallback; an explicit request value wins. Distilled two-stage uses fixed positive-only guidance. |
| `guidance_scale` | `float`, `None` | One-stage and ordinary Stage 1 common video/audio CFG fallback; an explicit request value wins. Distilled phases use fixed positive-only guidance. |
| `guidance_rescale` | `float`, `None` | Accepts only `None` or `0.0`; use the modality rescale fields. |
| `noise_scale` | `float`, `0.0` | Compatibility slot; LTX accepts only `0.0`. |
| `num_videos_per_prompt` | `int`, `1` | Output-count fallback; positive request `num_outputs_per_prompt` wins. |
| `generator` | generator or list, `None` | Explicit RNG; otherwise request generators/seeds are collated. Lists must match the effective output batch. |
| `latents` | tensor, `None` | Request tensors win. One-stage accepts packed `[B, S, C]` and validated unpacked video layouts. |
| `audio_latents` | tensor, `None` | One-stage initial audio latents; request tensors win and are collated. |
| `latents` | tensor, `None` | One-stage only. Request tensors win; packed `[B, S, C]` and validated unpacked video layouts are accepted. |
| `audio_latents` | tensor, `None` | One-stage only. Request tensors win and are collated. |
| `prompt_embeds` | tensor, `None` | Precomputed positive conditioning; requires `prompt_attention_mask` and cannot accompany `prompt`. |
| `negative_prompt_embeds` | tensor, `None` | One-stage precomputed negative conditioning; requires its mask and cannot accompany `negative_prompt`. |
| `negative_prompt_embeds` | tensor, `None` | One-stage and ordinary Stage 1 conditioning; requires its mask and cannot accompany `negative_prompt`. |
| `prompt_attention_mask` | tensor, `None` | Mask for positive embeddings; request `prompt_attention_mask` or `attention_mask` wins. |
| `negative_prompt_attention_mask` | tensor, `None` | Mask for negative embeddings; request negative mask fields win. |
| `decode_timestep` | float or list, `0.0` | Video-VAE decode timestep; request value wins. Lists may match 1, prompt batch, or output batch. |
@@ -273,13 +294,13 @@ request prompt payload; LTX guidance fields live in sampling `extra_args`.
### Recipe-Specific Request Capabilities
| Override | One-stage | Distilled two-stage |
|---|---|---|
| Guidance | Supported | Fixed positive-only |
| Negative prompt/embeddings | Supported | Rejected |
| `num_inference_steps` | Supported | Fixed at 8 for Stage 1; Stage 2 uses 3 |
| Custom `sigmas` | Supported | Rejected; both phases use fixed schedules |
| Video/audio latents | Supported | Rejected |
| Override | One-stage | Ordinary two-stage | Distilled two-stage |
|---|---|---|---|
| Guidance | Supported | Stage 1 only; Stage 2 is positive-only | Fixed positive-only |
| Negative prompt/embeddings | Supported | Supported by Stage 1 | Rejected |
| `num_inference_steps` | Supported | Controls Stage 1; Stage 2 uses 3 | Fixed at 8 for Stage 1; Stage 2 uses 3 |
| Custom `sigmas` | Supported | Rejected | Rejected; both phases use fixed schedules |
| Video/audio latents | Supported | Rejected | Rejected |
These capability checks apply equally to direct `forward` keywords and
values in `OmniDiffusionSamplingParams`; unsupported values fail instead of
@@ -309,6 +330,10 @@ bundled offline CLI do not currently expose `sigmas`.
`--max-num-seqs 1`.
- `--cfg-parallel-size` shards the complete LTX guidance plan, including STG,
modality guidance, and rescale-compatible prediction gathering.
- Sequence parallelism may pad audio latents. Pure Ulysses masks the padding;
Ring cannot, so audio length must be SP-divisible. Use `ring_degree=1` or a
divisible request shape.
- Cache-DiT is one-stage only; multi-stage configurations are rejected.
## Operational Notes
@@ -323,14 +348,15 @@ bundled offline CLI do not currently expose `sigmas`.
request parameter.
- For benchmarks, use `tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with
`tests/dfx/perf/scripts/run_diffusion_benchmark.py`.
- Current two-stage support is limited to the distilled checkpoint; official
LTX two-stage/HQ execution remains out of scope.
- Ordinary and full-distilled two-stage T2V/I2V are supported; HQ execution
remains out of scope.
## References
- <https://huggingface.co/Lightricks/LTX-2>
- <https://huggingface.co/Lightricks/LTX-2.3>
- <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
- <https://huggingface.co/diffusers/LTX-2.3-Distilled-Diffusers>
- [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md)
- [Diffusion execution modes](../../docs/user_guide/diffusion/execution_modes.md)
- [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md)
+152
View File
@@ -0,0 +1,152 @@
# dots.tts
> Offline continuous-AR TTS at 48 kHz (rednote-hilab)
## Summary
- Vendor: rednote-hilab
- Model: `rednote-hilab/dots.tts-soar`
- Task: Text-to-speech, zero-shot synthesis only
- Mode: Offline end-to-end example only (not yet wired for online serving —
see [Known limitations](#known-limitations))
- Maintainer: Community
## When to use this recipe
Use this recipe as a known-good starting point for running
`rednote-hilab/dots.tts-soar` offline on vLLM-Omni on consumer-class GPUs.
dots.tts is a ~1.7B-parameter continuous-AR TTS model (Qwen2.5-1.5B base LM
+ 344M DiT flow-matching head + 180M AudioVAE) that emits 48 kHz mono audio.
It follows the same "vLLM-native base LM + side-path computation" pattern as
VoxCPM2 — single-stage pipeline
`Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR
loopback → AudioVAE (streaming decode)` — with a plain Qwen2 backbone
instead of MiniCPM4, and no FSQ / residual-LM stage.
This is an early integration (see [Known limitations](#known-limitations)
below before depending on it for anything beyond offline experimentation).
## References
- Offline end-to-end script:
[`examples/offline_inference/text_to_speech/dots_tts/end2end.py`](../../examples/offline_inference/text_to_speech/dots_tts/end2end.py)
- Example guide:
[`examples/offline_inference/text_to_speech/README.md`](../../examples/offline_inference/text_to_speech/README.md#dotstts)
- Default deploy config:
[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
- Talker / pipeline source:
[`vllm_omni/model_executor/models/dots_tts/`](../../vllm_omni/model_executor/models/dots_tts/)
- Upstream: [rednote-hilab/dots.tts](https://github.com/rednote-hilab/dots.tts)
## Hardware Support
This recipe documents one tested 16 GB consumer-GPU configuration. Other
vendor sections (ROCm, NPU) and larger-VRAM configurations are welcome as
community validation lands.
## GPU
### 1 x RTX 5080 16GB (Single GPU, Minimum Recommended)
dots.tts (~1.7B params across the base LM + DiT + AudioVAE + CAM++
speaker encoder, bfloat16) fits comfortably on a single 16 GB GPU. The
bundled default config at
[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
(`gpu_memory_utilization: 0.8`, `max_num_seqs: 4`, `enforce_eager: true`,
`enable_prefix_caching: false`) loads cleanly with ~5.1 GiB for model
weights and ~0.3 GiB peak activation; the remainder of the configured
budget is available for KV cache. Total resident footprint at idle is
roughly **7-8 GiB / 16 GB** — the only tight spot in the full CUDA-Graph
roadmap would be step 8's graph capture (not implemented yet; this
release runs `enforce_eager: true`, so it doesn't apply today).
#### Environment
- OS: Linux (WSL2)
- Python: 3.12
- Driver / runtime: NVIDIA driver 595.95
- torch: 2.11.0+cu130
- vLLM: 0.26.0
- vLLM-Omni: 0.22.1.dev (current `main`)
#### Command
```bash
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
--model rednote-hilab/dots.tts-soar \
--text "Hello, this is a test of dots TTS running on vLLM Omni."
```
The deploy config at
[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
is loaded automatically by the model registry (HF `model_type=dots_tts`).
Pass `--deploy-config <path>` to override.
#### Verification
**T1 — offline zero-shot synthesis**:
```bash
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
--model rednote-hilab/dots.tts-soar \
--text "Hello, this is a test of dots TTS running on vLLM Omni."
```
Observed: `output_audio/output.wav`, 3.52 s @ 48 kHz mono. Single
one-shot process (init → one `generate()` → exit), so the reported
numbers include engine init and first-request warmup, not just steady-state
per-step throughput — same caveat as VoxCPM2's recipe. `Inference: 5.52s`,
`RTF: 1.569`.
Weight-loading breakdown from the same run (all tensors matched, no
missing/extra keys): 951/951 AudioVAE, 244/244 DiT, 270/270
patch_encoder, 198 Qwen2, 938/938 CAM++ speaker encoder.
Whisper transcription of the output matched the input text with no
dropped leading word (confirms the streaming-vocoder patch-boundary fix
described in [Known limitations](#known-limitations)).
#### Notes
- Output: 48 kHz mono WAV.
- Checkpoints: `rednote-hilab/dots.tts-soar` is the validated default
used throughout this recipe. `rednote-hilab/dots.tts-base` shares the
same architecture but is unvalidated in this repo. `rednote-hilab/dots.tts-mf`
(MeanFlow, 2-4 step) is not supported — see below.
- `enforce_eager: true` and `enable_prefix_caching: false` in the deploy
config are load-bearing, not just conservative defaults: with prefix
caching enabled, vLLM-Omni's prefix-cache multimodal-output merge path
does not preserve the `sparse_audio` marker this model relies on to
route audio output correctly, and generation silently truncates to a
single ~160 ms patch. Do not override `enable_prefix_caching` for this
model until that framework-level gap is fixed.
## Known limitations
- **No online serving yet.** dots.tts is not registered in
`vllm_omni/entrypoints/openai/tts_adapters/` (the framework other TTS
models in this repo use for `/v1/audio/speech`), so `vllm serve
rednote-hilab/dots.tts-soar --omni` will start but the OpenAI-compatible
speech endpoint does not know how to build request params for it yet.
Offline `Omni()` (as used in `end2end.py`) is the only supported path
today.
- **Voice cloning is not wired.** The CAM++ x-vector speaker encoder
weights load and are exercised by `load_weights()`, but nothing in the
prompt builder consumes a reference audio yet — generation is zero-shot
only. `end2end.py` has no `--ref-audio`/`--ref-text` flags.
- **`dots.tts-mf` (MeanFlow, 2-4 step) checkpoint is not supported.** Only
the fixed 10-step Euler DiT sampler used by `dots.tts-soar` /
`dots.tts-base` is implemented.
- **No CUDA graph capture.** The talker runs fully eager. voxcpm2's three
captured graphs (base LM decode, CFM solver, VAE decode) have no
dots.tts equivalent yet.
- **Concurrent requests do not scale.** Each request's 10-step DiT Euler
integration runs serially in the side path (no cross-request batching,
unlike voxcpm2's `enable_batched_cfm`). A community review of this
integration measured no throughput gain at `c=4` concurrent requests
versus `c=1`.
- **`SamplingParams.seed` does not control audio-generation randomness.**
The DiT's flow-matching noise is deterministically derived per-request
from a fixed internal seed (reproducible run-to-run, matching
voxcpm2's `deterministic_cfm_seed` convention) rather than from the
caller-supplied `seed` field — the same limitation voxcpm2 has today.
+46
View File
@@ -5,6 +5,7 @@ import sys
from pathlib import Path
import pytest
import yaml
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / ".buildkite" / "common" / "scripts"))
@@ -139,3 +140,48 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None:
assert step["plugins"][0]["kubernetes"]["podSpecPatch"]["imagePullSecrets"] == [
{"name": "swr-secret"},
]
COVERAGE_PILOT_LABELS = ("Diffusion · Bagel Test", "TTS · Qwen3-TTS Base Test")
MERGE_PIPELINE = Path(".buildkite/cuda/test-merge.yml")
def _surviving_labels(changed_files: list[str]) -> set[str]:
doc = yaml.safe_load(MERGE_PIPELINE.read_text(encoding="utf-8"))
rendered = _render_test_pipeline(doc, changed_files=changed_files)
labels: set[str] = set()
def walk(steps: list | None) -> None:
for step in steps or []:
if not isinstance(step, dict):
continue
if "label" in step:
labels.add(step["label"])
walk(step.get("steps"))
walk(rendered.get("steps"))
return labels
@pytest.mark.parametrize(
"changed_file",
[
".buildkite/common/scripts/run_cov_split.sh",
"pyproject.toml",
],
)
def test_coverage_pilots_selected_by_their_shared_inputs(changed_file: str) -> None:
"""The split helper and the coverage config both change what these jobs measure.
Without them in source_file_dependencies the jobs are filtered out of normal
PR builds, so a regression would only surface in a later nightly.
"""
labels = _surviving_labels([changed_file])
for pilot in COVERAGE_PILOT_LABELS:
assert pilot in labels, f"{pilot} dropped when {changed_file} changed"
def test_coverage_pilots_not_selected_by_unrelated_change() -> None:
labels = _surviving_labels(["vllm_omni/entrypoints/openai/serving_chat.py"])
for pilot in COVERAGE_PILOT_LABELS:
assert pilot not in labels
+4 -48
View File
@@ -6,7 +6,7 @@ Unit tests for StageConfigFactory and related classes.
import importlib
import warnings
from dataclasses import dataclass, fields
from dataclasses import dataclass
from pathlib import Path
from unittest.mock import patch
@@ -34,9 +34,8 @@ from vllm_omni.config.stage_config import (
load_deploy_config,
merge_pipeline_deploy,
pipeline_cfg_resolver,
strip_parent_engine_args,
)
from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys
from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -475,7 +474,7 @@ class TestStageConfigFactory:
cli_overrides = {
"gpu_memory_utilization": 0.9,
"model": "some_model", # Internal
"stage_configs_path": "/path", # Internal
"deploy_config": "/path", # Internal
"batch_timeout": 10, # Internal
}
@@ -483,7 +482,7 @@ class TestStageConfigFactory:
assert overrides["gpu_memory_utilization"] == 0.9
assert "model" not in overrides
assert "stage_configs_path" not in overrides
assert "deploy_config" not in overrides
assert "batch_timeout" not in overrides
def test_per_stage_override_excludes_internal_keys(self):
@@ -525,49 +524,6 @@ class TestStageResolutionHelpers:
assert "model" not in overrides
assert "parallel_config" not in overrides
def test_strip_parent_engine_args_reports_only_surprising_parent_overrides(self):
parent_fields = {f.name: f for f in fields(EngineArgs)}
filtered, overridden = strip_parent_engine_args(
{
"model": "some/model",
"stage_configs_path": "/tmp/stages.yaml",
"tensor_parallel_size": 4,
"worker_extension_cls": "some.Extension",
"custom_pipeline_args": {"pipeline_class": "demo.Pipeline"},
},
parent_fields=parent_fields,
keep_keys={"worker_extension_cls"},
strip_keys={"stage_configs_path"},
no_warn_keys={"model"},
)
assert filtered == {
"worker_extension_cls": "some.Extension",
"custom_pipeline_args": {"pipeline_class": "demo.Pipeline"},
}
assert overridden == ["tensor_parallel_size"]
def test_strip_parent_engine_args_keeps_allowed_media_access_controls(self):
parent_fields = {f.name: f for f in fields(EngineArgs)}
filtered, overridden = strip_parent_engine_args(
{
"model": "some/model",
"stage_configs_path": "/tmp/stages.yaml",
"allowed_local_media_path": "/data/qwentts",
"allowed_media_domains": ["example.com"],
},
parent_fields=parent_fields,
keep_keys={"allowed_local_media_path", "allowed_media_domains"},
strip_keys={"stage_configs_path"},
no_warn_keys={"model"},
)
assert filtered == {
"allowed_local_media_path": "/data/qwentts",
"allowed_media_domains": ["example.com"],
}
assert overridden == []
class TestPipelineDiscovery:
"""Tests for the central pipeline registry (``OMNI_PIPELINES``)."""
+1
View File
@@ -471,6 +471,7 @@ def test_sub_config_fields_match_structured_scopes():
"trust_remote_code",
"dtype",
"attention_backend",
"attention_config",
"moe_backend",
"hf_overrides",
"limit_mm_per_prompt",
@@ -60,9 +60,18 @@ def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
return torch.einsum("bhqk,bkhd->bqhd", probs, value)
def _cuda_flash_attn_usable() -> bool:
def _gpu_flash_attn_usable() -> bool:
if not torch.cuda.is_available():
return False
if torch.version.hip is not None:
for module in ("aiter", "flash_attn"):
try:
imported = __import__(module, fromlist=["flash_attn_varlen_func"])
if getattr(imported, "flash_attn_varlen_func", None) is not None:
return True
except ImportError:
pass
return False
try:
spec = find_spec("vllm.vllm_flash_attn")
if spec is None or spec.origin is None:
@@ -227,12 +236,11 @@ def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len,
assert st.adapter(POS).completed_chunks == before + (1 if commit_current else 0)
@pytest.mark.skipif(not _cuda_flash_attn_usable(), reason="usable CUDA FlashAttention is required")
@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
@pytest.mark.parametrize("history_chunks", [1, 3])
@pytest.mark.parametrize("action_len", [0, 3])
@pytest.mark.parametrize("commit_current", [False, True])
def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current):
pytest.importorskip("vllm.vllm_flash_attn")
torch.manual_seed(0)
device = torch.device("cuda")
dtype = torch.float16
+9 -4
View File
@@ -265,10 +265,15 @@ def test_fa_vs_sdpa():
print(f"Max absolute difference (valid region): {max_diff:.6f}")
print(f"Mean absolute difference (valid region): {mean_diff:.6f}")
# Assert that outputs are close
# Using higher tolerance for bfloat16 and different implementations
assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
# AITER's ROCm BF16 kernel is deterministic, but differs from ROCm SDPA by
# about 1.3% relative error. An absolute-only bound misclassifies that
# expected rounding difference for larger output values (for example,
# -2.65625 vs -2.6875). Keep the tighter historical bounds elsewhere.
if current_omni_platform.is_rocm():
torch.testing.assert_close(output_fa_valid, output_sdpa_valid, rtol=1e-2, atol=1.5e-2)
else:
assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
print("✓ Case 2 PASSED: FA and SDPA outputs are very close!")
@@ -416,8 +416,8 @@ def test_sp_correctness_advanced(model_name: str):
@pytest.mark.skipif(
not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
reason="Only tested on CUDA and XPU",
not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()),
reason="Only tested on CUDA, ROCM and XPU",
)
@pytest.mark.diffusion
@pytest.mark.parallel
@@ -447,8 +447,8 @@ def test_local_sp_padding_mask(monkeypatch: pytest.MonkeyPatch) -> None:
@pytest.mark.skipif(
not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
reason="Only tested on CUDA and XPU",
not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()),
reason="Only tested on CUDA, ROCM and XPU",
)
@pytest.mark.diffusion
@pytest.mark.parallel
@@ -468,8 +468,8 @@ def test_local_sp_padding_mask_no_padding(monkeypatch: pytest.MonkeyPatch) -> No
@pytest.mark.skipif(
not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
reason="Only tested on CUDA and XPU",
not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()),
reason="Only tested on CUDA, ROCM and XPU",
)
@pytest.mark.diffusion
@pytest.mark.parallel
@@ -0,0 +1,82 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import torch.nn.functional as F
from vllm.triton_utils import HAS_TRITON
pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion]
_HEAD_DIM = 128
_ROTARY_DIM = 96
_EPS = 1e-5
def _reference(q, k, q_weight, k_weight, rope_table):
q = F.rms_norm(q, (_HEAD_DIM,), q_weight, _EPS)
k = F.rms_norm(k, (_HEAD_DIM,), k_weight, _EPS)
half = _ROTARY_DIM // 2
cos = rope_table[..., :half].unsqueeze(1)
sin = rope_table[..., half:].unsqueeze(1)
def apply(x):
first = x[..., :half]
second = x[..., half:_ROTARY_DIM]
return torch.cat(
(
first * cos - second * sin,
second * cos + first * sin,
x[..., _ROTARY_DIM:],
),
dim=-1,
)
return apply(q), apply(k)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
@pytest.mark.parametrize("seq_len", [1, 257, 1024])
def test_fused_qk_norm_rope_matches_bf16_reference(seq_len):
from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
fused_qk_norm_rope,
)
torch.manual_seed(17)
heads = 14
qkv = torch.randn(
seq_len,
heads * _HEAD_DIM * 3,
device="cuda",
dtype=torch.bfloat16,
)
q = qkv[:, : heads * _HEAD_DIM].view(seq_len, heads, _HEAD_DIM)
k = qkv[:, heads * _HEAD_DIM : 2 * heads * _HEAD_DIM].view(
seq_len,
heads,
_HEAD_DIM,
)
q_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
k_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
freqs = torch.randn(seq_len, _ROTARY_DIM // 2, device="cuda")
rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1).to(torch.bfloat16)
expected_q, expected_k = _reference(
q,
k,
q_weight,
k_weight,
rope_table,
)
actual_q, actual_k = fused_qk_norm_rope(
q,
k,
q_weight,
k_weight,
rope_table,
_EPS,
)
torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)
@@ -24,11 +24,20 @@ from safetensors.torch import save_file
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServer, OmniServerParams
from vllm_omni.platforms import current_omni_platform
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
MODEL = "Tongyi-MAI/Z-Image-Turbo"
# ROCm's default FLASH_ATTN diffusion backend resolves to AITER. On a cold CI
# worker, JIT-building its MHA kernel can take longer than the server readiness
# timeout. This test validates LoRA switching rather than attention-backend or
# compiled-execution performance, so use the eager SDPA path on ROCm only.
ROCM_SERVER_ARGS = (
["--enforce-eager", "--diffusion-attention-backend", "TORCH_SDPA"] if current_omni_platform.is_rocm() else []
)
PROMPT = "a photo of a cat sitting on a laptop keyboard"
SIZE = "256x256"
@@ -42,7 +51,14 @@ server_params = [
server_args=[
"--num-gpus",
"1",
*ROCM_SERVER_ARGS,
],
# A cold-cache Z-Image load can take slightly more than the
# orchestrator's 900s default on CI (weights alone have been
# observed at ~690s and full orchestrator readiness at ~906s).
# Match the large-model offline runner's startup allowance without
# weakening timeout detection for every online-serving test.
init_timeout=1800,
),
id="zimage_turbo_lora",
),
@@ -23,7 +23,7 @@ def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1):
return pipeline
def test_prepare_video_latents_samples_directly_in_packed_token_space():
def test_prepare_video_latents_matches_official_values_and_token_major_layout():
pipeline = _make_pipeline(LTX2Pipeline)
pipeline.vae_spatial_compression_ratio = 8
pipeline.vae_temporal_compression_ratio = 8
@@ -44,6 +44,7 @@ def test_prepare_video_latents_samples_directly_in_packed_token_space():
)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
assert actual.stride()[1:] == (1, actual.shape[1])
def test_prepare_audio_latents_samples_directly_in_packed_token_space():
@@ -81,6 +82,42 @@ def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
assert original_num_frames == 1
assert padded_num_frames == 2
assert latents.shape == (1, 2, 128)
torch.testing.assert_close(latents[:, 1:], torch.zeros_like(latents[:, 1:]))
def test_prepare_audio_latents_request_rng_is_invariant_to_sp_padding():
pipeline_sp1 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=1)
pipeline_sp4 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
generator_sp1 = torch.Generator().manual_seed(42)
generator_sp4 = torch.Generator().manual_seed(42)
latents_sp1, _, _ = pipeline_sp1.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=8,
audio_latent_length=3,
dtype=torch.float32,
device=torch.device("cpu"),
generator=generator_sp1,
)
latents_sp4, _, _ = pipeline_sp4.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=8,
audio_latent_length=3,
dtype=torch.float32,
device=torch.device("cpu"),
generator=generator_sp4,
)
torch.testing.assert_close(latents_sp4[:, :3], latents_sp1, rtol=0, atol=0)
torch.testing.assert_close(latents_sp4[:, 3:], torch.zeros_like(latents_sp4[:, 3:]))
torch.testing.assert_close(
torch.randn(8, generator=generator_sp4),
torch.randn(8, generator=generator_sp1),
rtol=0,
atol=0,
)
def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents():
@@ -137,7 +174,9 @@ def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp():
assert original_num_frames == 10
assert padded_num_frames == 12
assert padded.shape == (1, 12, 8)
torch.testing.assert_close(padded, latent_ops.pack_audio_latents(latents))
packed = latent_ops.pack_audio_latents(latents)
torch.testing.assert_close(padded[:, :10], packed[:, :10])
torch.testing.assert_close(padded[:, 10:], torch.zeros_like(padded[:, 10:]))
def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference():
@@ -12,6 +12,19 @@ import torch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_velocity_from_x0_uses_official_host_scalar_sigma(mocker):
from vllm_omni.diffusion.models.ltx2.ltx2_guidance import velocity_from_x0
sigma = mocker.MagicMock()
sigma.to.return_value.item.return_value = 0.5
actual = velocity_from_x0(torch.tensor([1.0]), torch.tensor([0.0]), sigma)
sigma.to.assert_called_once_with(torch.float32)
sigma.to.return_value.item.assert_called_once_with()
torch.testing.assert_close(actual, torch.tensor([2.0]))
class TestCFGParallelHelpers:
"""Test LTX-2.3 CFG helper math without loading model weights."""
@@ -199,6 +212,7 @@ class TestCFGParallelHelpers:
),
attention_kwargs=None,
audio_scheduler=SimpleNamespace(sigmas=torch.stack([audio_sigma])),
original_audio_num_frames=state.audio.shape[1],
)
actual_video, actual_audio = executor.predict_parallel_guidance(
@@ -0,0 +1,483 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the fixed LTX refinement-phase adapter."""
from types import SimpleNamespace
import pytest
import torch
from safetensors.torch import save_file
from vllm_omni.diffusion.models.ltx2 import ltx2_phase_adapter
from vllm_omni.diffusion.models.ltx2.ltx2_adapter_parser import (
AdapterTarget,
parse_ltx_adapter,
)
from vllm_omni.diffusion.models.ltx2.ltx2_components import LTX23_TWO_STAGE_COMPONENT_PROFILE
from vllm_omni.diffusion.models.ltx2.ltx2_phase_adapter import LTXPhaseAdapterRuntime
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
LTX2_TWO_STAGE_RECIPE,
LTX_DISTILLED_ADAPTER_SLOT,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _target(
lora_a,
lora_b,
*,
source_module="proj",
module="proj",
target_slice=None,
):
return AdapterTarget(
source_module=source_module,
module=module,
slice=target_slice,
a_key="a",
b_key="b",
a_shape=tuple(lora_a.shape),
b_shape=tuple(lora_b.shape),
)
def _wrap(
layer,
lora_a,
lora_b,
*,
source_module="proj",
module="proj",
target_slice=None,
layer_fused=False,
):
target = _target(
lora_a,
lora_b,
source_module=source_module,
module=module,
target_slice=target_slice,
)
wrapper = ltx2_phase_adapter._PhaseAdapterLinear(
layer,
[target],
adapter_dtype=lora_a.dtype,
layer_fused=layer_fused,
)
wrapper.load_target(target, lora_a, lora_b)
wrapper.set_enabled(True)
return wrapper
class _LoRAAttention(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.to_qkv = torch.nn.Identity()
class _LoRATransformer(torch.nn.Module):
stacked_params_mapping = ((".attn1.to_qkv", ".attn1.to_q", "q"),)
def __init__(self) -> None:
super().__init__()
self.proj_in = torch.nn.Linear(4, 6, bias=False)
block = torch.nn.Module()
block.attn1 = _LoRAAttention()
self.transformer_blocks = torch.nn.ModuleList([block])
def test_ltx_adapter_parser_maps_official_and_packed_module_names(tmp_path):
adapter_path = tmp_path / "distilled.safetensors"
save_file(
{
"diffusion_model.text_embedding_projection.aggregate_embed.lora_A.weight": torch.randn(2, 4),
"diffusion_model.patchify_proj.lora_A.weight": torch.randn(2, 4),
"diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2),
"diffusion_model.transformer_blocks.0.attn1.to_q.lora_A.weight": torch.randn(2, 4),
"diffusion_model.transformer_blocks.0.attn1.to_q.lora_B.weight": torch.randn(4, 2),
},
adapter_path,
)
manifest = parse_ltx_adapter(_LoRATransformer(), str(adapter_path))
assert {(target.source_module, target.module, target.slice) for target in manifest.targets} == {
("proj_in", "proj_in", None),
("transformer_blocks.0.attn1.to_q", "transformer_blocks.0.attn1.to_qkv", "q"),
}
def test_ltx_adapter_parser_rejects_unpaired_b_tensor(tmp_path):
adapter_path = tmp_path / "distilled.safetensors"
save_file(
{"diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2)},
adapter_path,
)
with pytest.raises(ValueError, match="Missing paired LoRA A tensors"):
parse_ltx_adapter(_LoRATransformer(), str(adapter_path))
def _pipeline(*, dtype, quantization_config):
return SimpleNamespace(
pipeline_recipe=LTX2_TWO_STAGE_RECIPE,
component_profile=LTX23_TWO_STAGE_COMPONENT_PROFILE,
transformer=object(),
od_config=SimpleNamespace(
model="unused",
lora_path=None,
dtype=dtype,
quantization_config=quantization_config,
),
)
@pytest.mark.parametrize(
("dtype", "quantization_config", "layer_fused"),
[
(torch.float32, object(), False),
(torch.bfloat16, None, True),
],
)
def test_ltx_phase_adapter_selects_execution_from_quantization(
monkeypatch,
dtype,
quantization_config,
layer_fused,
):
pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config)
manifest = object()
captured = {}
def resolve_artifact(*_args, **kwargs):
assert not kwargs
return "adapter"
def parse(transformer, path):
assert transformer is pipeline.transformer
assert path == "adapter"
return manifest
class Runtime:
def __init__(self, transformer, parsed_manifest, **kwargs):
captured.update(transformer=transformer, manifest=parsed_manifest, **kwargs)
monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", resolve_artifact)
monkeypatch.setattr(ltx2_phase_adapter, "parse_ltx_adapter", parse)
monkeypatch.setattr(ltx2_phase_adapter, "LTXPhaseAdapterRuntime", Runtime)
runtime = ltx2_phase_adapter.build_ltx_phase_adapter(pipeline)
assert isinstance(runtime, Runtime)
assert captured == {
"transformer": pipeline.transformer,
"manifest": manifest,
"dtype": dtype,
"layer_fused": layer_fused,
}
def test_ltx_phase_adapter_rejects_non_bfloat16_layer_fused(monkeypatch):
pipeline = _pipeline(dtype=torch.float16, quantization_config=None)
monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", lambda *_args, **_kwargs: "adapter")
with pytest.raises(ValueError, match="bfloat16"):
ltx2_phase_adapter.build_ltx_phase_adapter(pipeline)
def test_ltx_phase_adapter_keeps_one_transformer_and_switches_a_fixed_slot(tmp_path):
class TinyTransformer(torch.nn.Module):
stacked_params_mapping = ()
def __init__(self) -> None:
super().__init__()
self.proj = torch.nn.Linear(2, 2, bias=False)
adapter_path = tmp_path / "distilled.safetensors"
save_file(
{
"diffusion_model.proj.lora_A.weight": torch.tensor([[1.0, 2.0]]),
"diffusion_model.proj.lora_B.weight": torch.tensor([[3.0], [4.0]]),
},
adapter_path,
)
transformer = TinyTransformer()
runtime = LTXPhaseAdapterRuntime(
transformer,
parse_ltx_adapter(transformer, str(adapter_path)),
dtype=torch.float32,
)
assert runtime.base_parameter_name("proj.weight") == "proj.base_layer.weight"
assert {name for name, _ in transformer.named_parameters()} == {"proj.base_layer.weight"}
with torch.no_grad():
transformer.proj.base_layer.weight.copy_(torch.eye(2))
runtime.finalize()
x = torch.tensor([[1.0, 1.0]])
runtime.activate(None)
torch.testing.assert_close(transformer.proj(x), torch.tensor([[1.0, 1.0]]))
runtime.activate(LTX_DISTILLED_ADAPTER_SLOT)
torch.testing.assert_close(transformer.proj(x), torch.tensor([[10.0, 13.0]]))
def test_ltx_layer_fused_adapter_matches_official_bf16_weight_fusion():
layer = torch.nn.Linear(4, 3, bias=True, dtype=torch.bfloat16)
base_weight = torch.tensor(
[
[0.5, -0.25, 0.125, 0.75],
[-0.375, 0.625, 0.25, -0.5],
[0.875, -0.125, -0.75, 0.375],
],
dtype=torch.bfloat16,
)
lora_a = torch.tensor(
[
[0.1015625, -0.203125, 0.3046875, -0.40625],
[0.5078125, -0.609375, 0.7109375, -0.8125],
],
dtype=torch.bfloat16,
)
lora_b = torch.tensor(
[[0.9140625, -0.1171875], [-0.21875, 0.3203125], [0.421875, -0.5234375]],
dtype=torch.bfloat16,
)
with torch.no_grad():
layer.weight.copy_(base_weight)
layer.bias.copy_(torch.tensor([0.125, -0.25, 0.5], dtype=torch.bfloat16))
wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
expected_weight = torch.matmul(lora_b, lora_a)
expected_weight.add_(base_weight)
x = torch.tensor([[0.6015625, -0.703125, 0.8046875, -0.90625]], dtype=torch.bfloat16)
expected = torch.nn.functional.linear(x, expected_weight, layer.bias)
assert torch.equal(wrapper(x), expected)
assert torch.equal(layer.weight, base_weight)
def test_ltx_layer_fused_row_parallel_uses_one_collective(monkeypatch):
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
class FakeRow(torch.nn.Module):
input_size = 4
output_size = 2
input_size_per_partition = 2
input_is_parallel = True
tp_rank = 0
tp_size = 2
reduce_results = True
skip_bias_add = False
return_bias = False
def __init__(self):
super().__init__()
self.weight = torch.nn.Parameter(torch.eye(2, dtype=torch.bfloat16))
self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16))
self.quant_method = UnquantizedLinearMethod()
monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", FakeRow)
reductions = []
def all_reduce(value):
reductions.append(value.clone())
return value + 1
monkeypatch.setattr(ltx2_phase_adapter, "tensor_model_parallel_all_reduce", all_reduce)
layer = FakeRow()
lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16)
lora_b = torch.tensor([[0.25], [-0.5]], dtype=torch.bfloat16)
wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
expected_weight = torch.matmul(lora_b, lora_a[:, :2])
expected_weight.add_(layer.weight)
x = torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16)
expected_local = torch.nn.functional.linear(x, expected_weight, layer.bias)
actual = wrapper(x)
assert len(reductions) == 1
assert torch.equal(reductions[0], expected_local)
assert torch.equal(actual, expected_local + 1)
def test_ltx_layer_fused_column_parallel_uses_local_weight_shard(monkeypatch):
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
class FakeColumn(torch.nn.Module):
input_size = 3
output_size = 4
output_size_per_partition = 2
tp_rank = 1
tp_size = 2
gather_output = False
skip_bias_add = False
return_bias = False
def __init__(self):
super().__init__()
self.weight = torch.nn.Parameter(
torch.tensor([[0.5, -0.25, 0.125], [-0.375, 0.625, 0.25]], dtype=torch.bfloat16)
)
self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16))
self.quant_method = UnquantizedLinearMethod()
monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn)
layer = FakeColumn()
base_weight = layer.weight.detach().clone()
lora_a = torch.tensor([[0.5, -0.25, 0.75]], dtype=torch.bfloat16)
lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16)
wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
expected_weight = torch.matmul(lora_b[2:], lora_a)
expected_weight.add_(base_weight)
x = torch.tensor([[1.0, -0.5, 0.25]], dtype=torch.bfloat16)
expected = torch.nn.functional.linear(x, expected_weight, layer.bias)
assert torch.equal(wrapper(x), expected)
assert torch.equal(layer.weight, base_weight)
def test_ltx_layer_fused_packed_qkv_updates_only_the_target_slice(monkeypatch):
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
class FakeColumn(torch.nn.Module):
pass
class FakeQKV(FakeColumn):
input_size = 4
head_size = 2
v_head_size = 2
total_num_heads = 2
total_num_kv_heads = 2
num_heads = 1
num_kv_heads = 1
num_kv_head_replicas = 1
tp_rank = 1
tp_size = 2
gather_output = False
skip_bias_add = True
return_bias = True
def __init__(self):
super().__init__()
self.weight = torch.nn.Parameter(
torch.arange(24, dtype=torch.float32).reshape(6, 4).to(torch.bfloat16) / 16
)
self.bias = torch.nn.Parameter(torch.arange(6, dtype=torch.float32).to(torch.bfloat16) / 8)
self.quant_method = UnquantizedLinearMethod()
monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn)
monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", FakeQKV)
layer = FakeQKV()
base_weight = layer.weight.detach().clone()
lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16)
lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16)
wrapper = _wrap(
layer,
lora_a,
lora_b,
source_module="to_k",
module="to_qkv",
target_slice="k",
layer_fused=True,
)
expected_weight = base_weight.clone()
expected_k = torch.matmul(lora_b[2:], lora_a)
expected_k.add_(base_weight[2:4])
expected_weight[2:4].copy_(expected_k)
x = torch.tensor([[1.0, -0.5, 0.25, -0.125]], dtype=torch.bfloat16)
expected_output = torch.nn.functional.linear(x, expected_weight)
actual_output, returned_bias = wrapper(x)
assert torch.equal(actual_output, expected_output)
assert returned_bias is layer.bias
assert torch.equal(layer.weight, base_weight)
@pytest.mark.parametrize(
("layer_factory", "target_slice", "expected_a", "expected_b", "expected_output"),
[
(
lambda: type(
"FakeColumn",
(torch.nn.Module,),
{
"input_size": 4,
"output_size": 4,
"output_size_per_partition": 2,
"tp_rank": 1,
"gather_output": False,
"forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)),
},
)(),
None,
torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
torch.tensor([[1.0], [2.0], [3.0], [4.0]]),
torch.tensor([[12963.0, 17284.0]]),
),
(
lambda: type(
"FakeRow",
(torch.nn.Module,),
{
"input_size": 4,
"output_size": 2,
"input_size_per_partition": 2,
"tp_rank": 1,
"reduce_results": False,
"forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)),
},
)(),
None,
torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
torch.tensor([[2.0], [3.0]]),
torch.tensor([[8600.0, 12900.0]]),
),
(
lambda: type(
"FakeQKV",
(torch.nn.Module,),
{
"input_size": 4,
"head_size": 2,
"v_head_size": 2,
"total_num_heads": 2,
"total_num_kv_heads": 2,
"num_heads": 1,
"num_kv_heads": 1,
"num_kv_head_replicas": 1,
"tp_rank": 1,
"forward": lambda self, x: x.new_zeros((*x.shape[:-1], 6)),
},
)(),
"k",
torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
torch.tensor([[1.0], [2.0], [3.0], [4.0]]),
torch.tensor([[0.0, 0.0, 12963.0, 17284.0, 0.0, 0.0]]),
),
],
)
def test_ltx_phase_adapter_uses_local_tp_lora_slices(
monkeypatch,
layer_factory,
target_slice,
expected_a,
expected_b,
expected_output,
):
layer = layer_factory()
layer.register_buffer("device_anchor", torch.empty(0))
layer_type = type(layer)
if target_slice == "k":
monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", layer_type)
elif "Column" in layer_type.__name__:
monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", layer_type)
else:
monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", layer_type)
wrapper = _wrap(layer, expected_a, expected_b, target_slice=target_slice)
torch.testing.assert_close(wrapper(torch.tensor([[1.0, 2.0, 3.0, 4.0]])), expected_output)
expected_local_a_width = getattr(layer, "input_size_per_partition", layer.input_size)
assert wrapper.pieces[0].lora_a.shape[-1] == expected_local_a_width
+362 -10
View File
@@ -13,13 +13,21 @@ import pytest
import torch
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.models.ltx2 import ltx2_components
from vllm_omni.diffusion.models.ltx2 import (
ltx2_components,
ltx2_transformer,
)
from vllm_omni.diffusion.models.ltx2.ltx2_components import (
LTX2_COMPONENT_PROFILE,
LTX2_DISTILLED_COMPONENT_PROFILE,
LTX2_TWO_STAGE_COMPONENT_PROFILE,
LTX23_COMPONENT_PROFILE,
LTX23_DISTILLED_COMPONENT_PROFILE,
LTX23_TWO_STAGE_COMPONENT_PROFILE,
_install_connector_attention,
detect_ltx_model_version,
resolve_ltx_artifact,
resolve_ltx_component_profile,
)
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import (
@@ -33,6 +41,7 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
LTXGuidancePlan,
LTXGuidanceSpec,
LTXModalityGuidance,
_repeat_batch,
build_perturbation_kwargs,
combine_guided_x0,
)
@@ -40,18 +49,29 @@ from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
LTX2_DISTILLED_TWO_STAGE_RECIPE,
LTX2_ONE_STAGE_RECIPE,
LTX2_TWO_STAGE_RECIPE,
LTX23_DISTILLED_TWO_STAGE_RECIPE,
LTX23_ONE_STAGE_RECIPE,
LTX23_TWO_STAGE_RECIPE,
LTX_DEFAULT_NEGATIVE_PROMPT,
LTX_POSITIVE_ONLY_RECIPE,
resolve_ltx_pipeline_recipe,
)
from vllm_omni.diffusion.models.ltx2.ltx2_request import (
LTXRequestInputs,
validate_ltx_checkpoint,
validate_pipeline_request,
)
from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs, validate_pipeline_request
from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2I2VDMD2Pipeline,
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
)
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
LTX2DistilledPipeline,
LTX2TwoStagePipeline,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -104,10 +124,22 @@ def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries():
)
assert issubclass(LTX2Pipeline, LTXRuntime)
assert issubclass(LTX2TwoStagePipeline, LTXRuntime)
assert issubclass(LTX2DistilledPipeline, LTXRuntime)
assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE
assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE
assert LTX2TwoStagePipeline.component_profile is LTX2_TWO_STAGE_COMPONENT_PROFILE
assert LTX2TwoStagePipeline.pipeline_recipe is LTX2_TWO_STAGE_RECIPE
assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE
for pipeline_cls, profile in (
(LTX2Pipeline, LTX2_COMPONENT_PROFILE),
(LTX2TwoStagePipeline, LTX2_TWO_STAGE_COMPONENT_PROFILE),
(LTX2DistilledPipeline, LTX2_DISTILLED_COMPONENT_PROFILE),
):
assert pipeline_cls._dit_modules == list(profile.dit_modules)
assert pipeline_cls._encoder_modules == list(profile.encoder_modules)
assert pipeline_cls._vae_modules == list(profile.vae_modules)
assert pipeline_cls._resident_modules == list(profile.resident_modules)
assert LTX2_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
assert LTX23_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
@@ -123,6 +155,72 @@ def test_ltx_checkpoint_version_detection_uses_metadata(tmp_path):
assert detect_ltx_model_version(str(tmp_path)) == "2"
def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch):
filename = "ltx-sidecar.safetensors"
calls = []
def fake_download(**kwargs):
calls.append(kwargs)
return "/cache/ltx-sidecar.safetensors"
monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
assert (
resolve_ltx_artifact(str(tmp_path / "model"), "Lightricks/test", filename) == "/cache/ltx-sidecar.safetensors"
)
assert calls == [{"repo_id": "Lightricks/test", "filename": filename}]
def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
filename = "ltx-sidecar.safetensors"
expected = tmp_path / filename
expected.write_bytes(b"sidecar")
monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup"))
assert resolve_ltx_artifact(str(tmp_path), "Lightricks/test", filename) == str(expected)
@pytest.mark.parametrize(
("expected_kind", "scheduler_config"),
[
("regular", {"use_dynamic_shifting": True, "shift_terminal": 0.1}),
("distilled", {"use_dynamic_shifting": False, "shift_terminal": None}),
],
)
def test_ltx_checkpoint_validation_accepts_matching_scheduler_metadata(expected_kind, scheduler_config):
validate_ltx_checkpoint(
scheduler_config,
expected_kind=expected_kind,
pipeline_name="TestLTXPipeline",
)
@pytest.mark.parametrize(
("expected_kind", "scheduler_config", "error"),
[
("regular", {"use_dynamic_shifting": False, "shift_terminal": None}, "regular non-distilled"),
("distilled", {"use_dynamic_shifting": True, "shift_terminal": 0.1}, "merged distilled"),
],
)
def test_ltx_checkpoint_validation_rejects_mismatched_scheduler_metadata(expected_kind, scheduler_config, error):
with pytest.raises(ValueError, match=error):
validate_ltx_checkpoint(
scheduler_config,
expected_kind=expected_kind,
pipeline_name="TestLTXPipeline",
)
def test_ltx_rejects_advanced_uaa_before_component_initialization():
od_config = SimpleNamespace(
model="unused",
parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa"),
)
with pytest.raises(ValueError, match="does not support ulysses_mode='advanced_uaa'"):
LTX2Pipeline(od_config=od_config)
def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
@@ -144,16 +242,104 @@ def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, m
assert pipe.reports_stage_durations
def test_ltx_one_stage_entry_exposes_batch_image_and_distributed_decode_contracts():
def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]}))
initialized = False
def stub_components(*_args, **_kwargs):
nonlocal initialized
initialized = True
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
with pytest.raises(ValueError, match="only one-stage LTX recipes"):
LTX2TwoStagePipeline(
od_config=SimpleNamespace(
model=str(tmp_path),
cache_backend="cache_dit",
)
)
assert not initialized
def test_ltx23_checkpoint_selects_version_specific_two_stage_profiles(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]}))
def stub_components(pipe, od_config):
pipe.od_config = od_config
pipe.vae_spatial_compression_ratio = 32
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda *_args: SimpleNamespace())
pipe = LTX2TwoStagePipeline(
od_config=SimpleNamespace(
model=str(tmp_path),
enable_diffusion_pipeline_profiler=False,
lora_path=None,
model_config={},
)
)
assert pipe.model_version == "2.3"
assert pipe.component_profile is LTX23_TWO_STAGE_COMPONENT_PROFILE
assert pipe.pipeline_recipe is LTX23_TWO_STAGE_RECIPE
def test_ltx23_checkpoint_selects_full_distilled_profile_and_recipe(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]}))
def stub_components(pipe, od_config):
pipe.od_config = od_config
pipe.vae_spatial_compression_ratio = 32
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
pipe = LTX2DistilledPipeline(
od_config=SimpleNamespace(
model=str(tmp_path),
enable_diffusion_pipeline_profiler=False,
)
)
assert pipe.model_version == "2.3"
assert pipe.component_profile is LTX23_DISTILLED_COMPONENT_PROFILE
assert pipe.pipeline_recipe is LTX23_DISTILLED_TWO_STAGE_RECIPE
assert pipe._phase_adapter is None
assert pipe.preserve_sp_padded_audio_duration
assert pipe.reports_stage_durations
def test_ltx_entries_expose_batch_image_and_distributed_decode_contracts():
assert LTX2Pipeline.supports_request_batch
assert LTX2Pipeline.support_image_input
assert LTX2Pipeline.unified_text_image_entry
assert LTX2Pipeline.distributed_video_decode
assert not LTX2TwoStagePipeline.supports_request_batch
assert LTX2TwoStagePipeline.support_image_input
assert LTX2TwoStagePipeline.unified_text_image_entry
assert not LTX2DistilledPipeline.supports_request_batch
assert LTX2DistilledPipeline.support_image_input
assert LTX2DistilledPipeline.unified_text_image_entry
def test_ltx_two_stage_resolves_request_embedded_images():
pipe = object.__new__(LTX2TwoStagePipeline)
request_inputs = SimpleNamespace(latents=None)
image = torch.zeros(3, 8, 8)
req = SimpleNamespace(prompts=[{"prompt": "image conditioned", "multi_modal_data": {"image": image}}])
assert pipe._resolve_request_image(req, None, request_inputs) is image
def test_ltx_dmd2_entries_retain_their_task_boundaries():
assert not LTX2T2VDMD2Pipeline.support_image_input
assert LTX2I2VDMD2Pipeline.support_image_input
@@ -330,6 +516,15 @@ def test_ltx_guidance_uses_official_close_to_disabled_semantics():
assert LTXGuidancePlan.build(LTXGuidanceSpec(video=guidance)).names == ("cond",)
def test_ltx_positive_only_guidance_preserves_token_major_layout():
tensor = torch.arange(24).view(1, 3, 8).transpose(1, 2).contiguous().transpose(1, 2)
repeated = _repeat_batch(tensor, 1)
assert repeated is tensor
assert repeated.stride() == tensor.stride()
def test_ltx_guidance_rescale_is_invariant_to_request_batching():
cond = torch.tensor(
[
@@ -380,6 +575,32 @@ def test_ltx_guidance_rescale_handles_zero_variance_prediction():
torch.testing.assert_close(actual, torch.zeros_like(cond))
def test_ltx_audio_guidance_rescale_ignores_sp_padding_tokens():
logical_cond = torch.tensor([[[1.0, 2.0], [3.0, 5.0]]])
logical_uncond = torch.tensor([[[0.5, -1.0], [2.0, 1.0]]])
guidance = LTXModalityGuidance(cfg_scale=3.0, rescale_scale=1.0)
expected = combine_guided_x0(
cond=logical_cond,
uncond_text=logical_uncond,
uncond_perturbed=0.0,
uncond_modality=0.0,
guidance=guidance,
)
padded_cond = torch.cat([logical_cond, torch.tensor([[[1000.0, -1000.0]]])], dim=1)
padded_uncond = torch.cat([logical_uncond, torch.tensor([[[-500.0, 500.0]]])], dim=1)
actual = combine_guided_x0(
cond=padded_cond,
uncond_text=padded_uncond,
uncond_perturbed=0.0,
uncond_modality=0.0,
guidance=guidance,
rescale_token_count=2,
)
torch.testing.assert_close(actual[:, :2], expected)
def test_ltx_guidance_builds_per_sample_stg_and_modality_masks():
plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance)
perturbations = build_perturbation_kwargs(plan, batch_size=2, reference=torch.ones(8, 1, 1))
@@ -547,11 +768,19 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
def test_ltx2_distilled_pipeline_shares_recipe_runtime():
assert issubclass(LTX2DistilledPipeline, LTXI2VConditioningMixin)
assert issubclass(LTX2DistilledPipeline, LTX2Pipeline)
assert not issubclass(LTX2DistilledPipeline, LTX2Pipeline)
assert LTX2DistilledPipeline._run_recipe is LTXRuntime._run_recipe
assert "_forward_request" not in LTX2DistilledPipeline.__dict__
def test_ltx_two_stage_entry_is_a_thin_pipeline_variant():
assert "_run_recipe" not in LTX2TwoStagePipeline.__dict__
assert "_forward_request" not in LTX2TwoStagePipeline.__dict__
assert "load_weights" not in LTX2TwoStagePipeline.__dict__
assert issubclass(LTX2TwoStagePipeline, LTXI2VConditioningMixin)
assert not issubclass(LTX2TwoStagePipeline, LTX2Pipeline)
def test_ltx_variants_share_denoise_loop_and_i2v_conditioning():
assert "_denoise_loop" not in LTX2Pipeline.__dict__
assert issubclass(LTX2Pipeline, LTXI2VConditioningMixin)
@@ -594,7 +823,8 @@ def test_denoise_executor_owns_progress_and_interrupt():
torch.testing.assert_close(state.audio, torch.tensor(11.0))
def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
request_inputs = LTXRequestInputs(
prompt="prompt",
negative_prompt="",
@@ -671,12 +901,15 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
prompt_context_sentinel = prompt_context
pipeline = object.__new__(LTX2DistilledPipeline)
pipeline = object.__new__(pipeline_cls)
torch.nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
pipeline.vae_spatial_compression_ratio = 32
pipeline.vae_temporal_compression_ratio = 8
pipeline.latent_upsampler = FakeUpsampler()
activated_slots = []
if pipeline_cls is LTX2TwoStagePipeline:
pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append)
object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
object.__setattr__(pipeline, "run_phase", run_phase)
object.__setattr__(pipeline, "decode_phase", decode_phase)
@@ -686,14 +919,18 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
assert len(phase_calls) == 2
assert phase_calls[1][2] is prompt_context_sentinel
if pipeline_cls is LTX2TwoStagePipeline:
assert activated_slots == [None, "ltx_distilled"]
torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0))
torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0))
expected_audio = 2.0 if pipeline_cls is LTX2TwoStagePipeline else 4.0
torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), expected_audio))
def test_ltx2_distilled_stage2_reapplies_final_resolution_i2v_conditioning():
@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
def test_ltx_two_stage_stage2_reapplies_final_resolution_i2v_conditioning(pipeline_cls):
from vllm_omni.diffusion.models.ltx2 import ltx2_latents
pipeline = object.__new__(LTX2DistilledPipeline)
pipeline = object.__new__(pipeline_cls)
torch.nn.Module.__init__(pipeline)
pipeline.vae_spatial_compression_ratio = 32
pipeline.vae_temporal_compression_ratio = 8
@@ -750,6 +987,18 @@ def test_ltx2_distilled_two_stage_recipe_is_fixed_positive_only():
assert (LTX2_DISTILLED_TWO_STAGE_RECIPE.height, LTX2_DISTILLED_TWO_STAGE_RECIPE.width) == (1024, 1536)
def test_ltx23_full_distilled_route_uses_merged_weights_without_lora():
assert resolve_ltx_component_profile("distilled_two_stage", "2.3") is LTX23_DISTILLED_COMPONENT_PROFILE
assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.3") is LTX23_DISTILLED_TWO_STAGE_RECIPE
assert LTX23_DISTILLED_TWO_STAGE_RECIPE is LTX2_DISTILLED_TWO_STAGE_RECIPE
assert all(phase.adapter_slot is None for phase in LTX23_DISTILLED_TWO_STAGE_RECIPE.phases)
assert LTX23_DISTILLED_COMPONENT_PROFILE.distilled_lora_filename is None
assert LTX23_DISTILLED_COMPONENT_PROFILE.vocoder_cls is LTX23_COMPONENT_PROFILE.vocoder_cls
assert LTX23_DISTILLED_COMPONENT_PROFILE.latent_upsampler_filename == (
"ltx-2.3-spatial-upscaler-x2-1.1.safetensors"
)
@pytest.mark.parametrize(
("request_inputs", "error"),
[
@@ -958,6 +1207,98 @@ def test_ltx_custom_sigmas_bypass_scheduler_shifting():
torch.testing.assert_close(timesteps, expected[:-1] * scheduler.config.num_train_timesteps)
@pytest.mark.parametrize(
("recipe", "steps", "stg_block"),
[
(LTX2_TWO_STAGE_RECIPE, 40, 29),
(LTX23_TWO_STAGE_RECIPE, 30, 28),
],
)
def test_ltx_official_two_stage_recipes_only_vary_by_model_defaults(recipe, steps, stg_block):
stage1, stage2 = recipe.phases
assert (recipe.width, recipe.height) == (1536, 1024)
assert recipe.num_inference_steps == steps
assert stage1.spatial_downscale == 2
assert stage1.adapter_slot is None
assert stage1.guidance.video.stg_blocks == (stg_block,)
assert stage1.guidance.audio.stg_blocks == (stg_block,)
assert stage2.input_transform == "spatial_upsample"
assert stage2.adapter_slot == "ltx_distilled"
assert stage2.guidance == LTXGuidanceSpec.positive_only()
assert stage2.num_inference_steps == 3
assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0)
# Official Stage 2 refines video only; its audio result is discarded.
assert (recipe.video_output_phase, recipe.audio_output_phase) == (1, 0)
def test_ltx_two_stage_entries_select_the_official_adapter_slots():
phase_switches = []
ordinary_pipeline = object.__new__(LTX2TwoStagePipeline)
ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append)
for phase in LTX2_TWO_STAGE_RECIPE.phases:
ordinary_pipeline._enter_phase(phase)
assert phase_switches == [None, "ltx_distilled"]
distilled_pipeline = object.__new__(LTX2DistilledPipeline)
for phase in LTX2_DISTILLED_TWO_STAGE_RECIPE.phases:
distilled_pipeline._enter_phase(phase)
def test_ltx23_gate_projection_declares_its_global_column_parallel_layout(monkeypatch):
class FakeColumn(torch.nn.Module):
def __init__(
self,
input_size,
output_size,
*,
bias,
gather_output,
return_bias,
quant_config,
prefix,
):
super().__init__()
self.input_size = input_size
self.output_size = output_size
self.bias_enabled = bias
self.gather_output = gather_output
self.return_bias = return_bias
self.quant_config = quant_config
self.prefix = prefix
class FakeRow(torch.nn.Module):
def __init__(self, *args, **kwargs):
super().__init__()
class FakeAttention(torch.nn.Module):
def __init__(self, **kwargs):
super().__init__()
monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumn)
monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRow)
monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeAttention)
monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2)
attention = ltx2_transformer.LTX2Attention(
query_dim=16,
heads=4,
dim_head=4,
qk_norm="rms_norm",
apply_gated_attention=True,
pack_qkv=False,
prefix="transformer_blocks.0.attn1",
)
gate = attention.to_gate_logits
assert isinstance(gate, FakeColumn)
assert gate.input_size == 16
assert gate.output_size == 4
assert gate.gather_output is False
assert gate.return_bias is False
assert gate.quant_config is None
assert gate.prefix == "transformer_blocks.0.attn1.to_gate_logits"
class TestLTXRequestParsing:
def test_request_resolves_official_and_explicit_negative_prompts_per_sample(self):
from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -1432,6 +1773,7 @@ class TestPipelineComponents:
def test_distilled_profile_keeps_upsampler_in_managed_resident_components(self):
assert LTX2_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler")
assert LTX23_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler")
class TestLTX23DecodeConditioning:
@@ -1477,6 +1819,11 @@ class TestRegistryIntegration:
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
assert _DIFFUSION_MODELS["LTX2Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX2Pipeline")
assert _DIFFUSION_MODELS["LTX2TwoStagePipeline"] == (
"ltx2",
"pipeline_ltx2_two_stage",
"LTX2TwoStagePipeline",
)
assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == (
"ltx2",
"pipeline_ltx2_two_stage",
@@ -1488,6 +1835,9 @@ class TestRegistryIntegration:
"LTX23ImageToVideoPipeline",
"LTX2TwoStagesPipeline",
"LTX2ImageToVideoTwoStagesPipeline",
"LTX2ImageToVideoTwoStagePipeline",
"LTX23TwoStagePipeline",
"LTX23ImageToVideoTwoStagePipeline",
}
assert removed_entries.isdisjoint(_DIFFUSION_MODELS)
@@ -1497,6 +1847,7 @@ class TestRegistryIntegration:
expected = [
"LTX2Pipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
@@ -1509,6 +1860,7 @@ class TestRegistryIntegration:
"model_class_name",
[
"LTX2Pipeline",
"LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
@@ -1,13 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import pytest
import torch
from torch import nn
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel, _make_rms_norm
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import (
LTX2AudioVideoAttnProcessor,
LTX2VideoTransformer3DModel,
_LTX2ParallelAttention,
_make_rms_norm,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -66,3 +73,83 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
assert video_timestep.expected_dims == audio_timestep.expected_dims == 2
assert not video_timestep.split_output
assert not audio_timestep.split_output
def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):
monkeypatch.setattr(LTX2AudioVideoAttnProcessor, "_is_sp_enabled", staticmethod(lambda: True))
processor = LTX2AudioVideoAttnProcessor()
hidden_states = torch.zeros(1, 3, 4)
encoder_hidden_states = torch.zeros(1, 4, 4)
key_padding_mask = torch.tensor([[True, True, True, False]])
actual = processor._prepare_attention_mask(
SimpleNamespace(),
hidden_states,
encoder_hidden_states,
key_padding_mask,
batch_size=1,
sequence_length=4,
)
torch.testing.assert_close(actual, key_padding_mask[:, None, None, :])
def _fake_ltx_parallel_attention():
calls = []
class Backend:
@staticmethod
def get_name():
return "FLASH_ATTN"
@staticmethod
def supports_attention_mask():
return True
def run(name):
def forward(query, _key, _value, _metadata):
calls.append(name)
return torch.zeros_like(query)
return forward
layer = object.__new__(_LTX2ParallelAttention)
nn.Module.__init__(layer)
layer.attn_backend = Backend
layer.attention = SimpleNamespace(forward=run("native"))
layer.sdpa_fallback = SimpleNamespace(forward=run("sdpa"))
layer.backend_pref = "FLASH_ATTN"
return layer, calls
@pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)])
def test_ltx_masked_cross_attention_uses_sdpa(query_length, key_length):
layer, calls = _fake_ltx_parallel_attention()
query = torch.zeros(1, query_length, 2, 4, dtype=torch.bfloat16)
key = value = torch.zeros(1, key_length, 2, 4, dtype=torch.bfloat16)
key_padding_mask = torch.tensor([[True] * (key_length - 1) + [False]])
metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
layer._run_local_attention(query, key, value, metadata)
assert calls == ["sdpa"]
def test_ltx_masked_self_attention_uses_sdpa():
layer, calls = _fake_ltx_parallel_attention()
query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
key_padding_mask = torch.tensor([[True, True, True, False]])
metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
layer._run_local_attention(query, key, value, metadata)
assert calls == ["sdpa"]
def test_ltx_unmasked_cross_attention_keeps_native_backend():
layer, calls = _fake_ltx_parallel_attention()
query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
layer._run_local_attention(query, key, value, None)
assert calls == ["native"]
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import sys
from multiprocessing.reduction import ForkingPickler
from types import SimpleNamespace
from unittest.mock import Mock
@@ -1428,34 +1429,219 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti
assert transcode_calls[0][1]["target_frame_count"] == 209
def test_ref2va_qwen_sampling_decodes_prepared_video_once(monkeypatch):
def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
decoded = np.arange(25 * 2 * 2 * 3, dtype=np.uint8).reshape(25, 2, 2, 3)
load_calls = []
decoded = np.arange(3 * 2 * 2 * 3, dtype=np.uint8).reshape(3, 2, 2, 3)
decode_calls = []
monkeypatch.setattr(
reference_video_module,
"_probe_video",
lambda _path: {"frame_count": len(decoded)},
lambda _path: {
"frame_count": 25,
"width": 2,
"height": 2,
},
)
monkeypatch.setattr(
reference_video_module,
"load_video_frames",
lambda path: load_calls.append(path) or decoded,
"_decode_video_frames_ffmpeg",
lambda path, **kwargs: decode_calls.append((path, kwargs)) 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,
):
assert decode_calls == [
(
"prepared.mp4",
{
"frame_count": 25,
"indices": [0, 12, 24],
"width": 2,
"height": 2,
},
)
]
for frame, expected in zip(sampled["frames"], decoded, strict=True):
np.testing.assert_array_equal(frame, expected)
assert sampled["block_timestamps"] == [0.25, 1.0]
@pytest.mark.parametrize(
("indices", "frame_count"),
[([0, 12], 25), (None, 2)],
)
def test_ref2va_ffmpeg_decode_reads_exact_rgb_frames(monkeypatch, indices, frame_count):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
expected = np.arange(2 * 2 * 3 * 3, dtype=np.uint8).reshape(2, 2, 3, 3)
observed = []
def fake_run(command, **kwargs):
observed.append((command, kwargs))
return SimpleNamespace(stdout=expected.tobytes())
monkeypatch.setattr(reference_video_module.subprocess, "run", fake_run)
actual = reference_video_module._decode_video_frames_ffmpeg(
"prepared.mp4",
frame_count=frame_count,
indices=indices,
width=3,
height=2,
)
np.testing.assert_array_equal(actual, expected)
command, kwargs = observed[0]
assert command[0] == "ffmpeg"
assert command[command.index("-frames:v") + 1] == "2"
if indices is None:
assert "-vf" not in command
else:
assert command[command.index("-vf") + 1] == "select=eq(n\\,0)+eq(n\\,12)"
assert command[-1] == "pipe:1"
assert kwargs == {"check": True, "capture_output": True}
def test_ref2va_load_video_frames_uses_ffmpeg_without_decord(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
expected = np.zeros((2, 2, 3, 3), dtype=np.uint8)
decode_calls = []
monkeypatch.setitem(sys.modules, "decord", None)
monkeypatch.setattr(
reference_video_module,
"_probe_video",
lambda _path: {
"frame_count": 2,
"width": 3,
"height": 2,
},
)
monkeypatch.setattr(
reference_video_module,
"_decode_video_frames_ffmpeg",
lambda path, **kwargs: decode_calls.append((path, kwargs)) or expected,
)
actual = reference_video_module.load_video_frames("prepared.mp4")
assert actual is expected
assert decode_calls == [
(
"prepared.mp4",
{
"frame_count": 2,
"width": 3,
"height": 2,
},
)
]
@pytest.mark.parametrize(
("kwargs", "message"),
[
({"frame_count": 0, "width": 3, "height": 2}, "video has no frames"),
({"frame_count": 1, "width": 0, "height": 2}, "invalid dimensions"),
(
{"frame_count": 1, "indices": [], "width": 3, "height": 2},
"invalid frame indices",
),
],
)
def test_ref2va_ffmpeg_decode_rejects_invalid_metadata(kwargs, message):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
from vllm_omni.errors import OmniClientError
with pytest.raises(OmniClientError, match=message):
reference_video_module._decode_video_frames_ffmpeg("prepared.mp4", **kwargs)
def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
calls = []
probe = {
"streams": [
{
"codec_type": "video",
"width": 4,
"height": 2,
"r_frame_rate": "24/1",
"nb_frames": "3",
"duration": "0.125",
"sample_aspect_ratio": "1:1",
}
],
"format": {"duration": "0.125", "size": "123"},
}
monkeypatch.setattr(
reference_video_module.subprocess,
"run",
lambda command, **kwargs: calls.append((command, kwargs)) or SimpleNamespace(stdout=json.dumps(probe)),
)
metadata = reference_video_module._probe_video("prepared.mp4")
assert metadata["frame_count"] == 3
assert len(calls) == 1
assert "-count_frames" not in calls[0][0]
def test_ref2va_probe_counts_frames_only_when_metadata_is_missing(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
calls = []
stream = {
"codec_type": "video",
"width": 4,
"height": 2,
"r_frame_rate": "24/1",
"duration": "0.125",
"sample_aspect_ratio": "1:1",
}
first_probe = {
"streams": [stream],
"format": {"duration": "0.125", "size": "123"},
}
second_probe = {
"streams": [{**stream, "nb_read_frames": "3"}],
"format": {"duration": "0.125", "size": "123"},
}
monkeypatch.setattr(
reference_video_module.subprocess,
"run",
lambda command, **kwargs: calls.append((command, kwargs))
or SimpleNamespace(stdout=json.dumps(first_probe if len(calls) == 1 else second_probe)),
)
metadata = reference_video_module._probe_video("prepared.mp4")
assert metadata["frame_count"] == 3
assert len(calls) == 2
assert "-count_frames" in calls[1][0]
def test_ref2va_ffmpeg_decode_rejects_incomplete_output(monkeypatch):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
from vllm_omni.errors import OmniClientError
monkeypatch.setattr(
reference_video_module.subprocess,
"run",
lambda *_args, **_kwargs: SimpleNamespace(stdout=b"short"),
)
with pytest.raises(OmniClientError, match="decoded 5 bytes.*expected 18"):
reference_video_module._decode_video_frames_ffmpeg(
"prepared.mp4",
frame_count=1,
width=3,
height=2,
)
def test_ref2va_two_video_recipe_tolerates_container_rounding(monkeypatch, tmp_path):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
@@ -0,0 +1,104 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
@pytest.mark.parametrize("num_q_heads,num_kv_heads,seq_len", [(8, 2, 16), (64, 8, 1)])
def test_qwen3_vl_npu_rope_matches_bnsd_reference(
num_q_heads: int,
num_kv_heads: int,
seq_len: int,
) -> None:
from vllm_omni.platforms.npu.models.minimax_h3 import _apply_rotary_pos_emb_npu
device = torch.device("npu")
dtype = torch.bfloat16
q = torch.randn(1, num_q_heads, seq_len, 128, device=device, dtype=dtype)
k = torch.randn(1, num_kv_heads, seq_len, 128, device=device, dtype=dtype)
cos = torch.randn(1, seq_len, 128, device=device, dtype=dtype)
sin = torch.randn(1, seq_len, 128, device=device, dtype=dtype)
q_out, k_out = _apply_rotary_pos_emb_npu(q, k, cos, sin)
cos_bnsd = cos.unsqueeze(1)
sin_bnsd = sin.unsqueeze(1)
def rotate_half(x: torch.Tensor) -> torch.Tensor:
x1, x2 = x.chunk(2, dim=-1)
return torch.cat((-x2, x1), dim=-1)
torch.testing.assert_close(
q_out,
q * cos_bnsd + rotate_half(q) * sin_bnsd,
atol=2e-2,
rtol=2e-2,
)
torch.testing.assert_close(
k_out,
k * cos_bnsd + rotate_half(k) * sin_bnsd,
atol=2e-2,
rtol=2e-2,
)
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
def test_qwen3_vl_npu_rope_preserves_batched_mrope_positions() -> None:
"""Keep per-sample Qwen3-VL M-RoPE frequencies on the BNSD NPU path."""
from types import SimpleNamespace
from vllm_omni.diffusion.models.minimax_h3.encoder import (
MiniMaxH3Qwen3VLTextRotaryEmbedding,
)
from vllm_omni.platforms.npu.models.minimax_h3 import _apply_rotary_pos_emb_npu
device = torch.device("npu")
dtype = torch.bfloat16
batch_size, seq_len, head_dim = 2, 16, 128
config = SimpleNamespace(
hidden_size=head_dim,
num_attention_heads=1,
rope_parameters={
"rope_theta": 1_000_000.0,
"mrope_section": [24, 20, 20],
},
)
rope = MiniMaxH3Qwen3VLTextRotaryEmbedding(config).to(device)
positions = torch.arange(seq_len, device=device)
position_ids = torch.stack(
(
torch.stack((positions, positions + 7)),
torch.stack((positions * 2, positions * 2 + 3)),
torch.stack((positions * 3, positions * 3 + 5)),
)
)
q = torch.randn(batch_size, 8, seq_len, head_dim, device=device, dtype=dtype)
k = torch.randn(batch_size, 2, seq_len, head_dim, device=device, dtype=dtype)
cos, sin = rope(q, position_ids)
q_out, k_out = _apply_rotary_pos_emb_npu(q, k, cos, sin)
cos_bnsd = cos.unsqueeze(1)
sin_bnsd = sin.unsqueeze(1)
def rotate_half(x: torch.Tensor) -> torch.Tensor:
x1, x2 = x.chunk(2, dim=-1)
return torch.cat((-x2, x1), dim=-1)
assert not torch.equal(cos[0], cos[1])
torch.testing.assert_close(
q_out,
q * cos_bnsd + rotate_half(q) * sin_bnsd,
atol=2e-2,
rtol=2e-2,
)
torch.testing.assert_close(
k_out,
k * cos_bnsd + rotate_half(k) * sin_bnsd,
atol=2e-2,
rtol=2e-2,
)
@@ -0,0 +1,46 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import torch.nn.functional as F
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
@pytest.mark.parametrize("shape", [(1, 16, 256), (2, 1, 256)])
def test_qwen3_vl_npu_swiglu_matches_split_reference(
shape: tuple[int, int, int],
) -> None:
from vllm_omni.platforms.npu.models.minimax_h3 import npu_swiglu_from_packed
packed = torch.randn(*shape, device="npu", dtype=torch.bfloat16)
gate, up = packed.chunk(2, dim=-1)
torch.testing.assert_close(
npu_swiglu_from_packed(packed),
F.silu(gate) * up,
atol=2e-2,
rtol=2e-2,
)
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
def test_qwen3_vl_npu_swiglu_matches_two_projection_reference() -> None:
from vllm_omni.platforms.npu.models.minimax_h3 import npu_swiglu_from_packed
x = torch.randn(1, 16, 128, device="npu", dtype=torch.bfloat16)
weight = torch.randn(256, 128, device="npu", dtype=torch.bfloat16)
packed = F.linear(x, weight)
gate = F.linear(x, weight[:128])
up = F.linear(x, weight[128:])
torch.testing.assert_close(
npu_swiglu_from_packed(packed),
F.silu(gate) * up,
atol=2e-2,
rtol=2e-2,
)
+32 -2
View File
@@ -36,15 +36,23 @@ class _ToyKernel:
# --- helper to build a toy module simulating online-FP8 post-load state ---
def _make_fp8_toy_module(out_features: int = 16, in_features: int = 32):
def _make_fp8_toy_module(
out_features: int = 16,
in_features: int = 32,
*,
quant_method_cls=None,
):
"""Return a module whose weight is ``qweight.t()`` (non-contiguous) and whose
``quant_method`` is an ``Fp8LinearMethod`` with a ``_ToyKernel``.
"""
from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
if quant_method_cls is None:
quant_method_cls = Fp8LinearMethod
# Create the quant-method instance without calling __init__ (avoids
# heavyweight upstream config dependencies). isinstance() still works.
qm = object.__new__(Fp8LinearMethod)
qm = object.__new__(quant_method_cls)
qm.fp8_linear = _ToyKernel()
# randn does not support float8 on CPU; zeros does and the test only
@@ -104,6 +112,28 @@ def test_rewrites_non_contiguous_weight():
assert module.weight.shape == (16, 32)
def test_rewrites_per_tensor_online_fp8_weight():
"""Regression test for vLLM's new online FP8 frontend (#45463)."""
from vllm.model_executor.layers.quantization.online.fp8 import (
Fp8PerTensorOnlineLinearMethod,
)
module = _make_fp8_toy_module(
16,
32,
quant_method_cls=Fp8PerTensorOnlineLinearMethod,
)
n_rewritten_layers = prepare_fp8_layers_for_fsdp(module)
assert n_rewritten_layers == 1
assert module.weight.is_contiguous()
assert module.weight.shape == (16, 32)
w, *_ = module.quant_method.fp8_linear._get_layer_params(module)
assert w.shape == (32, 16)
assert w.stride() == (1, 32)
def test_patched_get_layer_params_returns_transposed_view():
module = _make_fp8_toy_module(16, 32)
prepare_fp8_layers_for_fsdp(module)
@@ -189,7 +189,7 @@ def _generate_single_stage_video(
def _generate_bagel_image(
quantization_config: str | None = None,
diffusion_quantization_config: str | None = None,
num_inference_steps: int = 15,
) -> tuple[Any, float]:
"""Generate an image with BAGEL (multi-stage: LLM + Diffusion).
@@ -199,11 +199,11 @@ def _generate_bagel_image(
config_path = get_deploy_config_path("ci/bagel.yaml")
omni_kwargs: dict[str, Any] = {
"model": "ByteDance-Seed/BAGEL-7B-MoT",
"stage_configs_path": config_path,
"deploy_config": config_path,
"stage_init_timeout": 300,
}
if quantization_config:
omni_kwargs["quantization_config"] = quantization_config
if diffusion_quantization_config:
omni_kwargs["diffusion_quantization_config"] = diffusion_quantization_config
model_name = omni_kwargs.pop("model")
with OmniRunner(model_name, **omni_kwargs) as runner:
@@ -380,21 +380,21 @@ def test_single_stage_ltx2_fp8_uses_less_memory():
@hardware_test(res={"cuda": "H100"})
def test_bagel_fp8_generates_image():
"""BAGEL with FP8 quantization_config generates a valid image.
"""BAGEL with diffusion-stage FP8 generates a valid image.
FP8 should only apply to the diffusion stage (Stage-1), not the
LLM stage (Stage-0). We verify this by checking:
1. Image is generated successfully
2. LLM stage finish_reason is 'stop' (not 'length' from garbled output)
"""
image, _ = _generate_bagel_image(quantization_config="fp8")
image, _ = _generate_bagel_image(diffusion_quantization_config="fp8")
image.save("test_bagel_fp8.png")
@hardware_test(res={"cuda": "H100"})
def test_bagel_bf16_generates_image():
"""BAGEL without quantization generates a valid image (baseline)."""
image, _ = _generate_bagel_image(quantization_config=None)
image, _ = _generate_bagel_image(diffusion_quantization_config=None)
image.save("test_bagel_bf16.png")
@@ -95,6 +95,7 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata(
"resolution": 512,
"width": None,
"height": None,
"num_inference_steps": 1,
}
@@ -314,7 +314,7 @@ def test_bagel_text2img_mooncake_connector(run_level):
temp_config_file = _resolve_deploy_config(temp_config_file, run_level)
with OmniRunner(
"ByteDance-Seed/BAGEL-7B-MoT",
stage_configs_path=temp_config_file,
deploy_config=temp_config_file,
stage_init_timeout=300,
) as runner:
generated_image = _generate_bagel_image(runner.omni)
@@ -251,7 +251,7 @@ def test_bagel_img2img_shared_memory_connector(run_level):
config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level)
with OmniRunner(
"ByteDance-Seed/BAGEL-7B-MoT",
stage_configs_path=config_path,
deploy_config=config_path,
) as runner:
generated_image = _generate_bagel_img2img(runner.omni, input_image)
if run_level == "advanced_model":
@@ -266,7 +266,7 @@ def test_bagel_text2img_shared_memory_connector(run_level):
config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level)
with OmniRunner(
"ByteDance-Seed/BAGEL-7B-MoT",
stage_configs_path=config_path,
deploy_config=config_path,
) as runner:
generated_image = _generate_bagel_text2img(runner.omni)
if run_level == "advanced_model":
+117 -50
View File
@@ -20,6 +20,11 @@ from PIL import Image
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", choices=("official", "omni"), required=True)
parser.add_argument(
"--pipeline-kind",
choices=("one_stage", "distilled", "two_stage"),
default="one_stage",
)
parser.add_argument("--request", type=Path, required=True)
parser.add_argument("--output-dir", type=Path, required=True)
parser.add_argument("--model")
@@ -27,6 +32,8 @@ def _parse_args() -> argparse.Namespace:
parser.add_argument("--official-root", type=Path)
parser.add_argument("--checkpoint", type=Path)
parser.add_argument("--gemma-root", type=Path)
parser.add_argument("--spatial-upsampler", type=Path)
parser.add_argument("--distilled-lora", type=Path)
parser.add_argument("--enable-layerwise-offload", action="store_true")
return parser.parse_args()
@@ -91,9 +98,13 @@ def _configure_official_sdpa(pipeline: Any) -> None:
attention=attention,
masked_attention=attention,
)
owners_and_attributes = (
(pipeline.stage, "_transformer_builder"),
owners_and_attributes = [
(pipeline.prompt_encoder, "_embeddings_processor_builder"),
]
owners_and_attributes.extend(
(stage, "_transformer_builder")
for name in ("stage", "stage_1", "stage_2")
if (stage := getattr(pipeline, name, None)) is not None
)
for owner, attribute in owners_and_attributes:
builder = getattr(owner, attribute)
@@ -104,6 +115,12 @@ def _configure_official_sdpa(pipeline: Any) -> None:
)
def _require_path(path: Path | None, description: str) -> str:
if path is None or not path.is_file():
raise ValueError(f"Official {description} is required and must exist: {path}")
return str(path)
@torch.inference_mode()
def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
if args.official_root is None or args.checkpoint is None or args.gemma_root is None:
@@ -113,16 +130,54 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
torch.backends.cuda.enable_cudnn_sdp(False)
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_pipelines.distilled import DistilledPipeline
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
from ltx_pipelines.utils.args import ImageConditioningInput
from ltx_pipelines.utils.types import OffloadMode
pipeline = TI2VidOneStagePipeline(
checkpoint_path=str(args.checkpoint),
gemma_root=str(args.gemma_root),
loras=(),
offload_mode=OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE,
)
if request.get("pipeline_kind", "one_stage") != args.pipeline_kind:
raise ValueError(
f"Request pipeline kind {request.get('pipeline_kind')!r} does not match {args.pipeline_kind!r}"
)
checkpoint = _require_path(args.checkpoint, "checkpoint")
gemma_root = str(args.gemma_root)
if not Path(gemma_root).is_dir():
raise ValueError(f"Official Gemma root is required and must be a directory: {gemma_root}")
offload_mode = OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE
pipeline: Any
if args.pipeline_kind == "one_stage":
pipeline = TI2VidOneStagePipeline(
checkpoint_path=checkpoint,
gemma_root=gemma_root,
loras=(),
offload_mode=offload_mode,
)
elif args.pipeline_kind == "distilled":
pipeline = DistilledPipeline(
distilled_checkpoint_path=checkpoint,
spatial_upsampler_path=_require_path(args.spatial_upsampler, "spatial upsampler"),
gemma_root=gemma_root,
loras=(),
offload_mode=offload_mode,
)
else:
distilled_lora = [
LoraPathStrengthAndSDOps(
path=_require_path(args.distilled_lora, "distilled LoRA"),
strength=1.0,
sd_ops=LTXV_LORA_COMFY_RENAMING_MAP,
)
]
pipeline = TI2VidTwoStagesPipeline(
checkpoint_path=checkpoint,
distilled_lora=distilled_lora,
spatial_upsampler_path=_require_path(args.spatial_upsampler, "spatial upsampler"),
gemma_root=gemma_root,
loras=(),
offload_mode=offload_mode,
)
_configure_official_sdpa(pipeline)
image_path = request.get("image")
images = (
@@ -140,34 +195,45 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
)
]
)
video, audio = pipeline(
prompt=request["prompt"],
negative_prompt=request["negative_prompt"],
seed=request["seed"],
height=request["height"],
width=request["width"],
num_frames=request["num_frames"],
frame_rate=request["fps"],
num_inference_steps=request["num_inference_steps"],
video_guider_params=MultiModalGuiderParams(
cfg_scale=request["video_cfg_scale"],
stg_scale=request["video_stg_scale"],
rescale_scale=request["video_rescale_scale"],
modality_scale=request["video_modality_scale"],
skip_step=0,
stg_blocks=request["video_stg_blocks"],
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=request["audio_cfg_scale"],
stg_scale=request["audio_stg_scale"],
rescale_scale=request["audio_rescale_scale"],
modality_scale=request["audio_modality_scale"],
skip_step=0,
stg_blocks=request["audio_stg_blocks"],
),
images=images,
max_batch_size=4,
)
if args.pipeline_kind == "distilled":
video, audio = pipeline(
prompt=request["prompt"],
seed=request["seed"],
height=request["height"],
width=request["width"],
num_frames=request["num_frames"],
frame_rate=request["fps"],
images=images,
)
else:
video, audio = pipeline(
prompt=request["prompt"],
negative_prompt=request["negative_prompt"],
seed=request["seed"],
height=request["height"],
width=request["width"],
num_frames=request["num_frames"],
frame_rate=request["fps"],
num_inference_steps=request["num_inference_steps"],
video_guider_params=MultiModalGuiderParams(
cfg_scale=request["video_cfg_scale"],
stg_scale=request["video_stg_scale"],
rescale_scale=request["video_rescale_scale"],
modality_scale=request["video_modality_scale"],
skip_step=0,
stg_blocks=request["video_stg_blocks"],
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=request["audio_cfg_scale"],
stg_scale=request["audio_stg_scale"],
rescale_scale=request["audio_rescale_scale"],
modality_scale=request["audio_modality_scale"],
skip_step=0,
stg_blocks=request["audio_stg_blocks"],
),
images=images,
max_batch_size=4,
)
video_tensor = torch.cat([chunk.detach().cpu() for chunk in video], dim=0)
_save_outputs(
args.output_dir,
@@ -179,6 +245,7 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
"attention_backend": "torch_sdpa",
"official_revision": os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION"),
"checkpoint": str(args.checkpoint),
"pipeline_kind": args.pipeline_kind,
},
)
@@ -193,15 +260,9 @@ def _unwrap_omni_output(output: Any) -> tuple[Any, Any, int]:
multimodal_output = frames.multimodal_output or {}
audio = multimodal_output.get("audio")
audio_sample_rate = multimodal_output.get("audio_sample_rate")
if frames.is_pipeline_output and isinstance(frames.request_output, OmniRequestOutput):
frames = frames.request_output
multimodal_output = frames.multimodal_output or {}
audio = multimodal_output.get("audio", audio)
audio_sample_rate = multimodal_output.get("audio_sample_rate", audio_sample_rate)
if isinstance(frames, OmniRequestOutput):
if not frames.images:
raise ValueError("No video frames found in OmniRequestOutput")
frames = frames.images
if not frames.images:
raise ValueError("No video frames found in OmniRequestOutput")
frames = frames.images
if isinstance(frames, list) and len(frames) == 1:
frames = frames[0]
@@ -253,6 +314,10 @@ def _canonical_video(video: Any) -> torch.Tensor:
def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
if args.model is None or args.model_class_name is None:
raise ValueError("Omni backend requires --model and --model-class-name")
if request.get("pipeline_kind", "one_stage") != args.pipeline_kind:
raise ValueError(
f"Request pipeline kind {request.get('pipeline_kind')!r} does not match {args.pipeline_kind!r}"
)
attention_config = {"default": {"backend": "TORCH_SDPA"}}
@@ -273,7 +338,8 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
parallel_config=DiffusionParallelConfig(),
)
try:
model_class_name = get_model_class_name(omni)
detected_model_class_name = get_model_class_name(omni)
model_class_name = args.model_class_name
sampling_params = OmniDiffusionSamplingParams(
height=request["height"],
width=request["width"],
@@ -289,10 +355,9 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
key: value for key, value in request.items() if key.startswith("video_") or key.startswith("audio_")
}
apply_declared_extra_args(sampling_params, get_extra_body_params(model_class_name), guidance)
prompt: dict[str, Any] = {
"prompt": request["prompt"],
"negative_prompt": request["negative_prompt"],
}
prompt: dict[str, Any] = {"prompt": request["prompt"]}
if request.get("negative_prompt"):
prompt["negative_prompt"] = request["negative_prompt"]
image_path = request.get("image")
if image_path is not None:
with Image.open(str(image_path)) as source_image:
@@ -312,6 +377,8 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
"attention_backend": "torch_sdpa",
"model": args.model,
"model_class_name": model_class_name,
"model_config_class_name": detected_model_class_name,
"pipeline_kind": args.pipeline_kind,
},
)
finally:
@@ -3,10 +3,9 @@
"""E2E accuracy guard against a pinned Lightricks LTX pipeline revision.
The comparison runs both runtimes through PyTorch SDPA and uses
``max_batch_size=4`` in the official reference to match Omni's fused guidance
batch. Video and audio guidance use the official non-HQ one-stage defaults;
only the generation shape and step count are reduced for CI runtime.
The original reduced one-stage guards remain unchanged. Four complementary
default-shape cases cover both model versions, both two-stage weight families,
and T2V/I2V without expanding every combination into a duplicate golden.
"""
from __future__ import annotations
@@ -17,8 +16,9 @@ import shutil
import subprocess
import sys
import time
from dataclasses import dataclass
from dataclasses import asdict, dataclass, replace
from pathlib import Path
from typing import Literal
import numpy as np
import pytest
@@ -53,74 +53,260 @@ NEGATIVE_PROMPT = (
# Both runtimes use PyTorch SDPA with the current Torch dispatch defaults.
ATTENTION_BACKEND = "torch_sdpa"
VIDEO_SSIM_MEAN_THRESHOLD = 0.95
VIDEO_SSIM_MIN_THRESHOLD = 0.90
VIDEO_PSNR_MEAN_THRESHOLD = 30.0
AUDIO_RELATIVE_L2_THRESHOLD = 0.2
AUDIO_COSINE_THRESHOLD = 0.95
@dataclass(frozen=True)
class LTXArtifact:
repo_id: str
filename: str
revision: str
env: str
repo_type: Literal["model", "dataset"] = "model"
@dataclass(frozen=True)
class LTXAccuracyThresholds:
video_ssim_mean: float
video_ssim_min: float
video_psnr_mean_db: float
# Zero disables gating for waveform-sensitive audio that is not expected
# to align with the official output.
audio_relative_l2: float
audio_cosine_similarity: float
STRICT_THRESHOLDS = LTXAccuracyThresholds(
video_ssim_mean=0.95,
video_ssim_min=0.90,
video_psnr_mean_db=30.0,
audio_relative_l2=0.20,
audio_cosine_similarity=0.95,
)
@dataclass(frozen=True)
class LTXAccuracyCase:
name: str
pipeline_kind: Literal["one_stage", "distilled", "two_stage"]
model_id: str
model_revision: str
model_env: str
model_class_name: str
checkpoint_repo: str
checkpoint_filename: str
checkpoint_revision: str
checkpoint_env: str
stg_block: int
checkpoint: LTXArtifact
width: int
height: int
num_frames: int
num_inference_steps: int
seed: int
stg_block: int | None
thresholds: LTXAccuracyThresholds
prompt: str = PROMPT
image_repo: str | None = None
image_filename: str | None = None
image_revision: str | None = None
gemma_model_id: str | None = None
gemma_model_revision: str | None = None
gemma_model_env: str | None = None
image: LTXArtifact | None = None
spatial_upsampler: LTXArtifact | None = None
distilled_lora: LTXArtifact | None = None
# Official uses block streaming; Omni uses layerwise offload. Both preserve
# bf16 arithmetic while keeping full-resolution two-stage cases on one H100.
enable_layerwise_offload: bool = False
CASES = (
LTX2_REVISION = "47da56e2ad66ce4125a9922b4a8826bf407f9d0a"
LTX23_REVISION = "4229404625088d21c4f112eb640fb04a0900ee25"
LTX2_CHECKPOINT = LTXArtifact(
repo_id="Lightricks/LTX-2",
filename="ltx-2-19b-dev.safetensors",
revision=LTX2_REVISION,
env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT",
)
LTX2_DISTILLED_CHECKPOINT = replace(
LTX2_CHECKPOINT,
filename="ltx-2-19b-distilled.safetensors",
env="VLLM_TEST_LTX2_DISTILLED_OFFICIAL_CHECKPOINT",
)
LTX2_UPSAMPLER = replace(
LTX2_CHECKPOINT,
filename="ltx-2-spatial-upscaler-x2-1.0.safetensors",
env="VLLM_TEST_LTX2_UPSAMPLER",
)
LTX2_DISTILLED_LORA = replace(
LTX2_CHECKPOINT,
filename="ltx-2-19b-distilled-lora-384.safetensors",
env="VLLM_TEST_LTX2_DISTILLED_LORA",
)
LTX23_CHECKPOINT = LTXArtifact(
repo_id="Lightricks/LTX-2.3",
filename="ltx-2.3-22b-dev.safetensors",
revision=LTX23_REVISION,
env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
)
LTX23_DISTILLED_CHECKPOINT = replace(
LTX23_CHECKPOINT,
# The pinned Diffusers checkpoint contains the original merged distilled
# Transformer. Keep the reference Transformer on the same version; the
# spatial upsampler is versioned independently and remains on 1.1 below.
filename="ltx-2.3-22b-distilled.safetensors",
env="VLLM_TEST_LTX23_DISTILLED_OFFICIAL_CHECKPOINT",
)
LTX23_UPSAMPLER = replace(
LTX23_CHECKPOINT,
filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
env="VLLM_TEST_LTX23_UPSAMPLER",
)
LTX23_DISTILLED_LORA = replace(
LTX23_CHECKPOINT,
filename="ltx-2.3-22b-distilled-lora-384-1.1.safetensors",
env="VLLM_TEST_LTX23_DISTILLED_LORA",
)
I2V_IMAGE = LTXArtifact(
repo_id="huggingface/documentation-images",
filename="diffusers/svd/rocket.png",
revision="645d8364f0c7a101180b364811b5a11a362e4010",
env="VLLM_TEST_LTX_I2V_IMAGE",
repo_type="dataset",
)
LEGACY_CASES = (
LTXAccuracyCase(
name="ltx2",
pipeline_kind="one_stage",
model_id="Lightricks/LTX-2",
model_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a",
model_revision=LTX2_REVISION,
model_env="VLLM_TEST_LTX2_MODEL",
model_class_name="LTX2Pipeline",
checkpoint_repo="Lightricks/LTX-2",
checkpoint_filename="ltx-2-19b-dev.safetensors",
checkpoint_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a",
checkpoint_env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT",
checkpoint=LTX2_CHECKPOINT,
width=512,
height=384,
num_frames=25,
num_inference_steps=20,
seed=42,
stg_block=29,
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx2_3",
pipeline_kind="one_stage",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
model_class_name="LTX2Pipeline",
checkpoint_repo="Lightricks/LTX-2.3",
checkpoint_filename="ltx-2.3-22b-dev.safetensors",
checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25",
checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
checkpoint=LTX23_CHECKPOINT,
width=512,
height=384,
num_frames=25,
num_inference_steps=20,
seed=42,
stg_block=28,
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx2_3_i2v",
pipeline_kind="one_stage",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
model_class_name="LTX2Pipeline",
checkpoint_repo="Lightricks/LTX-2.3",
checkpoint_filename="ltx-2.3-22b-dev.safetensors",
checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25",
checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
checkpoint=LTX23_CHECKPOINT,
width=512,
height=384,
num_frames=25,
num_inference_steps=20,
seed=42,
stg_block=28,
image_repo="huggingface/documentation-images",
image_filename="diffusers/svd/rocket.png",
image_revision="645d8364f0c7a101180b364811b5a11a362e4010",
thresholds=STRICT_THRESHOLDS,
image=I2V_IMAGE,
),
)
TWO_STAGE_CASES = (
LTXAccuracyCase(
name="ltx2_distilled_t2v",
pipeline_kind="distilled",
model_id="rootonchair/LTX-2-19b-distilled",
model_revision="388e2846f54aae51687498ffb6b27c7c2c9ce9e5",
model_env="VLLM_TEST_LTX2_DISTILLED_MODEL",
model_class_name="LTX2DistilledPipeline",
checkpoint=LTX2_DISTILLED_CHECKPOINT,
spatial_upsampler=LTX2_UPSAMPLER,
width=1536,
height=1024,
num_frames=121,
num_inference_steps=8,
seed=10,
stg_block=None,
gemma_model_id="Lightricks/LTX-2",
gemma_model_revision=LTX2_REVISION,
gemma_model_env="VLLM_TEST_LTX2_MODEL",
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx23_distilled_i2v",
pipeline_kind="distilled",
model_id="diffusers/LTX-2.3-Distilled-Diffusers",
model_revision="432e0d3c2d1769aaa4d295f9243f7062bf6b47ee",
model_env="VLLM_TEST_LTX23_DISTILLED_MODEL",
model_class_name="LTX2DistilledPipeline",
checkpoint=LTX23_DISTILLED_CHECKPOINT,
spatial_upsampler=LTX23_UPSAMPLER,
width=1536,
height=1024,
num_frames=121,
num_inference_steps=8,
seed=10,
stg_block=None,
gemma_model_id="diffusers/LTX-2.3-Diffusers",
gemma_model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
gemma_model_env="VLLM_TEST_LTX23_MODEL",
thresholds=STRICT_THRESHOLDS,
image=I2V_IMAGE,
),
LTXAccuracyCase(
name="ltx2_two_stage_layer_fused_t2v",
pipeline_kind="two_stage",
model_id="Lightricks/LTX-2",
model_revision=LTX2_REVISION,
model_env="VLLM_TEST_LTX2_MODEL",
model_class_name="LTX2TwoStagePipeline",
checkpoint=LTX2_CHECKPOINT,
spatial_upsampler=LTX2_UPSAMPLER,
distilled_lora=LTX2_DISTILLED_LORA,
width=1536,
height=1024,
num_frames=121,
num_inference_steps=40,
seed=10,
stg_block=29,
enable_layerwise_offload=True,
thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx23_two_stage_layer_fused_i2v",
pipeline_kind="two_stage",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
model_class_name="LTX2TwoStagePipeline",
checkpoint=LTX23_CHECKPOINT,
spatial_upsampler=LTX23_UPSAMPLER,
distilled_lora=LTX23_DISTILLED_LORA,
width=1536,
height=1024,
num_frames=121,
num_inference_steps=30,
seed=10,
stg_block=28,
enable_layerwise_offload=True,
thresholds=STRICT_THRESHOLDS,
image=I2V_IMAGE,
),
)
CASES = (*LEGACY_CASES, *TWO_STAGE_CASES)
def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None:
start = time.perf_counter()
subprocess.run(command, env=env, timeout=timeout, check=True)
@@ -209,6 +395,7 @@ def _resolve_model(case: LTXAccuracyCase) -> Path:
"model_index.json",
"audio_vae/*",
"connectors/*",
"latent_upsampler/*",
"processor/*",
"scheduler/*",
"text_encoder/config.json",
@@ -223,74 +410,82 @@ def _resolve_model(case: LTXAccuracyCase) -> Path:
)
def _resolve_gemma_root(model: Path) -> Path:
def _resolve_gemma_root(case: LTXAccuracyCase, model: Path) -> Path:
configured_root = os.environ.get("VLLM_TEST_LTX_GEMMA_ROOT")
if configured_root:
root = Path(configured_root)
assert root.is_dir(), f"Gemma root not found: {root}"
return root
if case.gemma_model_id is not None:
configured_model = os.environ.get(case.gemma_model_env or "")
if configured_model and Path(configured_model).is_dir():
return Path(configured_model)
return Path(
snapshot_download(
repo_id=case.gemma_model_id,
revision=case.gemma_model_revision,
allow_patterns=[
"processor/*",
"text_encoder/*",
"tokenizer/*",
],
)
)
return model
def _resolve_checkpoint(case: LTXAccuracyCase, model: Path) -> Path:
configured_checkpoint = os.environ.get(case.checkpoint_env)
if configured_checkpoint:
checkpoint = Path(configured_checkpoint)
assert checkpoint.is_file(), f"Official checkpoint not found: {checkpoint}"
return checkpoint
model_checkpoint = model / case.checkpoint_filename
if model_checkpoint.is_file():
return model_checkpoint
def _resolve_artifact(artifact: LTXArtifact, model: Path | None = None) -> Path:
configured_path = os.environ.get(artifact.env)
if configured_path:
path = Path(configured_path)
assert path.is_file(), f"Configured LTX artifact not found: {path}"
return path
if model is not None:
model_path = model / artifact.filename
if model_path.is_file():
return model_path
return Path(
hf_hub_download(
repo_id=case.checkpoint_repo,
filename=case.checkpoint_filename,
revision=case.checkpoint_revision,
repo_id=artifact.repo_id,
repo_type=None if artifact.repo_type == "model" else artifact.repo_type,
filename=artifact.filename,
revision=artifact.revision,
)
)
def _resolve_image(case: LTXAccuracyCase) -> Path | None:
if case.image_filename is None:
if case.image is None:
return None
if case.image_repo is None or case.image_revision is None:
raise ValueError(f"Incomplete image source for LTX accuracy case {case.name!r}.")
configured_image = os.environ.get("VLLM_TEST_LTX_I2V_IMAGE")
if configured_image:
image = Path(configured_image)
assert image.is_file(), f"LTX I2V conditioning image not found: {image}"
return image
return Path(
hf_hub_download(
repo_id=case.image_repo,
repo_type="dataset",
filename=case.image_filename,
revision=case.image_revision,
)
)
return _resolve_artifact(case.image)
def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
request: dict[str, object] = {
"pipeline_kind": case.pipeline_kind,
"prompt": case.prompt,
"negative_prompt": NEGATIVE_PROMPT,
"width": 512,
"height": 384,
"num_frames": 25,
"negative_prompt": "" if case.pipeline_kind == "distilled" else NEGATIVE_PROMPT,
"width": case.width,
"height": case.height,
"num_frames": case.num_frames,
"fps": 24,
"num_inference_steps": 20,
"seed": 42,
"video_cfg_scale": 3.0,
"audio_cfg_scale": 7.0,
"video_stg_scale": 1.0,
"audio_stg_scale": 1.0,
"video_modality_scale": 3.0,
"audio_modality_scale": 3.0,
"video_rescale_scale": 0.7,
"audio_rescale_scale": 0.7,
"video_stg_blocks": [case.stg_block],
"audio_stg_blocks": [case.stg_block],
"num_inference_steps": case.num_inference_steps,
"seed": case.seed,
}
if case.pipeline_kind != "distilled":
assert case.stg_block is not None
request.update(
video_cfg_scale=3.0,
audio_cfg_scale=7.0,
video_stg_scale=1.0,
audio_stg_scale=1.0,
video_modality_scale=3.0,
audio_modality_scale=3.0,
video_rescale_scale=0.7,
audio_rescale_scale=0.7,
video_stg_blocks=[case.stg_block],
audio_stg_blocks=[case.stg_block],
)
if image is not None:
request["image"] = str(image.resolve())
return request
@@ -299,20 +494,30 @@ def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
assert reference.shape == prediction.shape
assert reference.ndim == 4 and reference.shape[-1] == 3
device = torch.device("cpu")
ssim = StructuralSimilarityIndexMeasure(data_range=1.0).to(device)
psnr = PeakSignalNoiseRatio(data_range=1.0).to(device)
ssim_scores: list[float] = []
psnr_scores: list[float] = []
for reference_frame, prediction_frame in zip(reference, prediction, strict=True):
reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0)
prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0)
ssim_scores.append(float(StructuralSimilarityIndexMeasure(data_range=1.0)(prediction_tensor, reference_tensor)))
psnr_scores.append(float(PeakSignalNoiseRatio(data_range=1.0)(prediction_tensor, reference_tensor)))
difference = np.abs(reference.astype(np.float64) - prediction.astype(np.float64))
max_abs = 0.0
absolute_error_sum = 0.0
with torch.inference_mode():
for reference_frame, prediction_frame in zip(reference, prediction, strict=True):
reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0).to(device)
prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0).to(device)
ssim_scores.append(float(ssim(prediction_tensor, reference_tensor)))
psnr_scores.append(float(psnr(prediction_tensor, reference_tensor)))
ssim.reset()
psnr.reset()
difference = np.abs(reference_frame.astype(np.float64) - prediction_frame.astype(np.float64))
max_abs = max(max_abs, float(difference.max()))
absolute_error_sum += float(difference.sum())
return {
"ssim_mean": float(np.mean(ssim_scores)),
"ssim_min": float(np.min(ssim_scores)),
"psnr_mean_db": float(np.mean(psnr_scores)),
"max_abs": float(difference.max()),
"mean_abs": float(difference.mean()),
"max_abs": max_abs,
"mean_abs": absolute_error_sum / reference.size,
}
@@ -343,13 +548,19 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
@pytest.mark.benchmark
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None:
"""Compare official and Omni raw AV outputs from the same E2E request."""
def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None:
"""Compare official and Omni raw AV outputs from the same single-H100 request."""
configured_artifact_root = os.environ.get("VLLM_TEST_LTX_ARTIFACT_ROOT")
if configured_artifact_root:
accuracy_artifact_root = Path(configured_artifact_root)
accuracy_artifact_root.mkdir(parents=True, exist_ok=True)
output_root = reset_artifact_dir(accuracy_artifact_root / "ltx_official" / case.name)
official_root, official_revision = _official_source(accuracy_artifact_root / "ltx_official")
model = _resolve_model(case)
gemma_root = _resolve_gemma_root(model)
checkpoint = _resolve_checkpoint(case, model)
gemma_root = _resolve_gemma_root(case, model)
checkpoint = _resolve_artifact(case.checkpoint, model)
spatial_upsampler = _resolve_artifact(case.spatial_upsampler, model) if case.spatial_upsampler is not None else None
distilled_lora = _resolve_artifact(case.distilled_lora, model) if case.distilled_lora is not None else None
image = _resolve_image(case)
request_path = output_root / "request.json"
request_path.write_text(json.dumps(_request(case, image), indent=2) + "\n")
@@ -359,8 +570,13 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
str(runner),
"--request",
str(request_path),
"--pipeline-kind",
case.pipeline_kind,
]
if os.environ.get("VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", "").lower() in {"1", "true", "yes", "on"}:
enable_layerwise_offload = case.enable_layerwise_offload or os.environ.get(
"VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", ""
).lower() in {"1", "true", "yes", "on"}
if enable_layerwise_offload:
runner_args.append("--enable-layerwise-offload")
env = os.environ.copy()
env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision
@@ -372,10 +588,9 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
)
official_output = output_root / "official"
_run(
_official_runner_prefix()
+ runner_args
+ [
official_args = _official_runner_prefix() + runner_args
official_args.extend(
[
"--backend",
"official",
"--output-dir",
@@ -386,12 +601,16 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
str(checkpoint),
"--gemma-root",
str(gemma_root),
],
env=env,
]
)
if spatial_upsampler is not None:
official_args.extend(["--spatial-upsampler", str(spatial_upsampler)])
if distilled_lora is not None:
official_args.extend(["--distilled-lora", str(distilled_lora)])
_run(official_args, env=env)
omni_output = output_root / "omni"
_run(
omni_args = (
[sys.executable]
+ runner_args
+ [
@@ -403,9 +622,9 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
str(model),
"--model-class-name",
case.model_class_name,
],
env=env,
]
)
_run(omni_args, env=env)
official_metadata = json.loads((official_output / "metadata.json").read_text())
omni_metadata = json.loads((omni_output / "metadata.json").read_text())
@@ -423,18 +642,27 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
result = {
"case": case.name,
"task": "i2v" if image is not None else "t2v",
"pipeline_kind": case.pipeline_kind,
"model_class_name": case.model_class_name,
"attention_backend": ATTENTION_BACKEND,
"official_revision": official_revision,
"model_revision": case.model_revision,
"checkpoint_revision": case.checkpoint_revision,
"checkpoint_revision": case.checkpoint.revision,
"spatial_upsampler_revision": (case.spatial_upsampler.revision if case.spatial_upsampler is not None else None),
"distilled_lora_revision": case.distilled_lora.revision if case.distilled_lora is not None else None,
"enable_layerwise_offload": enable_layerwise_offload,
"thresholds": asdict(case.thresholds),
"request": _request(case, image),
"video": video_metrics,
"audio": audio_metrics,
}
(output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(result, indent=2))
assert video_metrics["ssim_mean"] >= VIDEO_SSIM_MEAN_THRESHOLD
assert video_metrics["ssim_min"] >= VIDEO_SSIM_MIN_THRESHOLD
assert video_metrics["psnr_mean_db"] >= VIDEO_PSNR_MEAN_THRESHOLD
assert audio_metrics["relative_l2"] <= AUDIO_RELATIVE_L2_THRESHOLD
assert audio_metrics["cosine_similarity"] >= AUDIO_COSINE_THRESHOLD
assert video_metrics["ssim_mean"] >= case.thresholds.video_ssim_mean
assert video_metrics["ssim_min"] >= case.thresholds.video_ssim_min
assert video_metrics["psnr_mean_db"] >= case.thresholds.video_psnr_mean_db
if case.thresholds.audio_relative_l2 > 0:
assert audio_metrics["relative_l2"] <= case.thresholds.audio_relative_l2
if case.thresholds.audio_cosine_similarity > 0:
assert audio_metrics["cosine_similarity"] >= case.thresholds.audio_cosine_similarity
@@ -17,7 +17,7 @@ from vllm_omni.platforms import current_omni_platform
# Override in CI if needed: AUDIOX_TEST_MODEL=<model-or-local-path>
models = [os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random")]
# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
_OMNI_RUNNER_PARAMS = [(m, None, {"model_class_name": "AudioXPipeline"}) for m in models]
pytestmark = [
@@ -43,7 +43,7 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
BAGEL_STAGE_CONFIG = get_deploy_config_path("ci/bagel.yaml")
DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
# (model, stage_configs_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
# (model, deploy_config_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
_OMNI_RUNNER_PARAM = (MODEL, BAGEL_STAGE_CONFIG)
@@ -126,7 +126,7 @@ def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[s
]
# (model, stage_configs_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
# (model, deploy_config_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
_cosy_deployment = get_deploy_config_path("cosyvoice3.yaml")
_cosy_model_path = str(_resolve_model_dir())
_OMNI_RUNNER_PARAMS = [
@@ -17,7 +17,7 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
MODEL = "baidu/ERNIE-Image"
# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
_OMNI_RUNNER_PARAM = (
MODEL,
None,
@@ -96,7 +96,7 @@ def test_offline_voice_clone_zh(async_chunk: bool) -> None:
with OmniRunner(
MODEL,
stage_configs_path=_get_deploy_config(async_chunk=async_chunk),
deploy_config=_get_deploy_config(async_chunk=async_chunk),
stage_init_timeout=600,
trust_remote_code=True,
) as omni_runner:
@@ -52,7 +52,7 @@ pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
def omni() -> Generator[Omni, None, None]:
with OmniRunner(
MODEL_NAME,
stage_configs_path=str(STAGE_CONFIG_PATH),
deploy_config=str(STAGE_CONFIG_PATH),
trust_remote_code=True,
) as runner:
yield runner.omni
@@ -84,7 +84,6 @@ def longcat_omni_runner() -> Generator[OmniRunner, None, None]:
with OmniRunner(
model,
seed=42,
stage_configs_path=None,
model_class_name="LongCatVideoAvatarPipeline",
additional_config={
"model_type": "avatar-v1.5",
@@ -38,7 +38,7 @@ _AR_PATCH_SIZE = 16
MODEL_PATH = "bytedance-research/MammothModa2-Preview"
T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
_OMNI_RUNNER_PARAM = (MODEL_PATH, None, {"deploy_config": T2I_DEPLOY_CONFIG})
_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_CONFIG)
# Golden pixel reference file. Set UPDATE_GOLDEN=1 to regenerate.
_GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"
@@ -17,7 +17,7 @@ models = ["openbmb/MiniCPM-o-4_5"]
_CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
test_params = [(model, None, {"deploy_config": _CI_DEPLOY, "trust_remote_code": True}) for model in models]
test_params = [(model, _CI_DEPLOY, {"trust_remote_code": True}) for model in models]
def get_question(prompt_type: str = "text") -> str:
Loaded 100 of 186 files, more files were not shown because too many files have changed in this diff. Show more