Compare commits

...
Author SHA1 Message Date
Hongsheng Liu 089851e55a Merge origin/main into codex-vae-spatial-decode-auto
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 23:49:11 +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
liwenerandamy-why-3459 877beeaa49 [Perf][MiniCPM-O-4.5]Add CUDA Graph For CFM DiT estimator (#6082)
Signed-off-by: liwener <liwener.c@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-13 09:36:17 +08:00
Hongsheng Liu 8b8d3f56db [Misc] Consolidate LingBot text-to-video runner (#6076)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 08:42:37 +08:00
Hongsheng Liu c57178c0e1 [Dependency] Upgrade Cache-DiT to 1.5.0 (#6065)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 08:41:59 +08:00
mjZhaoElaineandClaude Opus 5 9a59475015 [CI/Build] Fix ReadTheDocs build under RTD's seeded pip 23.1 (#6129)
Signed-off-by: Mengjie Zhao <zmj0129@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-08-13 06:52:01 +08:00
Nick Cao 12f8168179 [Tests] Enable tiny model testing for Qwen-Image Edit and EditPlus (#5656)
Signed-off-by: Nick Cao <ncao@redhat.com>
2026-08-12 16:03:49 -04:00
Zhenghan Song 72a02b492a [Bugfix] Fix diffusion TTS adapter lookup in speech serving (#6121)
Signed-off-by: HaningZS <resossr0v0@gmail.com>
2026-08-12 15:20:12 -04:00
Juan Pablo ZuluagaandYueqian Lin 5358f5db01 [Perf][TTS] omnivoice hot path: D2H batching, mask caching, cached text embeddings (#5174)
Signed-off-by: JuanPZuluaga <juanz9312@gmail.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-08-12 17:15:35 +00:00
MaciejBalaNV e677360168 Add cosmos-guardrail dependency and update error message (#6107)
Signed-off-by: Maciej Bala <mbala@nvidia.com>
2026-08-12 17:04:07 +00:00
Juan Pablo ZuluagaandYueqian Lin 208c6776a8 [Perf][TTS] glm_tts: hoist loop-invariant text embed, RoPE, mask and CFG batch out of Euler loop (#5068)
Signed-off-by: JuanPZuluaga <juanz9312@gmail.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-08-12 17:02:02 +00:00
Hongsheng Liu 3466ff9b98 docs: standardize VAE Parallel naming
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 23:06:06 +08:00
fcf1085395 [Refactor][OutputProcessor 2/8]: OmniRequestOutput should inherit RequestOutput and no nested wrap-up (#5146)
Signed-off-by: Boao Shi <aoibosh@connect.hku.hk>
Signed-off-by: bowieshi <aoibosh@connect.hku.hk>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: R2-Y <ruiruyang2@gmail.com>
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Signed-off-by: Shi Boao Bowie <aoibosh@connect.hku.hk>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: ruirui(rein) yang <73573651+R2-Y@users.noreply.github.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 22:58:23 +08:00
Hongsheng Liu e16473c6b5 docs: align VAE parallel file convention
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 22:24:58 +08:00
Hongsheng Liu 14d67bcd60 docs: align VAE parallelism file names
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 22:15:51 +08:00
Hongsheng Liu 5234e3599c docs: generalize VAE parallelism design
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 22:07:54 +08:00
Wenjia LiandHongsheng Liu 38a2500bdf [Perf] MiniMax H3 Qwen3-VL support fused RMSNorm on NPU (#5915)
Signed-off-by: Wenjia Li <wjialish@gmail.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 19:12:54 +08:00
amy-why-3459 7cd826d28c [Bugfix][MiniCPM-o] Align Daily-Omni offline loading and duplex soft-… (#6095)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-12 18:36:31 +08:00
wangyuandHongsheng Liu 31384cb39c [CI/Build] Add nightly DockerHub publish and cleanup to release pipeline (#6048)
Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: [Your Name] [Your Email]
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 17:30:55 +08:00
Yuekai ZhangandClaude Fable 5 5be7541d01 [Model] Support NVIDIA-NemotronLabs-VoiceChat-11B offline speech-to-speech (#5842)
Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-08-12 16:25:48 +08:00
9b0df3909a [Feature] timestamps for TTS via a forced-aligner pooling (#4795)
Signed-off-by: wjinxu <1299461899@qq.com>
Co-authored-by: xurui <xurui1@modelbest.onaliyun.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-08-12 05:33:34 +00:00
Hongsheng Liu 3aefb31727 docs(diffusion): remove VAE policy attribution
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 12:01:32 +08:00
Hongsheng Liu c8210b8b5c fix(diffusion): address Wan auto decode review feedback
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 11:49:47 +08:00
Hongsheng Liu 879d19c4c2 Merge branch 'main' into codex-vae-spatial-decode-auto 2026-08-12 11:38:12 +08:00
Hongsheng Liu a15b5b4bcd [CI/Build] Add examples policy to PR skills (#6046)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 11:37:34 +08:00
Hongsheng Liu 3c3989ed04 docs: align quantization overview with navigation (#6074)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 11:37:01 +08:00
Rahul SteigerandRahul Steiger 854bff7fe5 [Bugfix] Fix Wan spatial reshard boundary (#6062)
Signed-off-by: Rahul Steiger <rsteiger@nvl72d106-T18.cm.cluster>
Co-authored-by: Rahul Steiger <rsteiger@nvl72d106-T18.cm.cluster>
2026-08-12 11:36:24 +08:00
amy-why-3459 b792578d89 [CI][MiniCPM-o] Add MiniCPM-o 4.5 perf coverage to the ready gate (#6079)
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-08-12 11:32:36 +08:00
zzh bdb1d996be [Bugfix] Fix HunyuanImage3 accuracy test (#5981)
Signed-off-by: zzh <943967662@qq.com>
2026-08-12 02:40:05 +00:00
Hongsheng Liu af1825e07d feat(diffusion): select Wan VAE decode mode automatically
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-11 17:36:35 +08:00
263 changed files with 17705 additions and 2411 deletions

No files matched your search

+5 -5
View File
@@ -345,7 +345,7 @@ steps:
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
@@ -361,7 +361,7 @@ steps:
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
@@ -377,7 +377,7 @@ steps:
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
@@ -393,7 +393,7 @@ steps:
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
@@ -409,7 +409,7 @@ steps:
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.3.0
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
+12
View File
@@ -227,6 +227,18 @@ steps:
"
mirror_hardwares: h100_1
- label: "Omni · MiniCPM-o 4.5 Perf Test"
timeout_in_minutes: 60
source_file_dependencies:
- tests/dfx/perf/tests/test_minicpmo_4_5_ready.json
- vllm_omni/deploy/minicpmo_4_5.yaml
- vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and cuda'
mirror_hardwares: h100_1
- label: "TTS · VoxCPM2 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_voxcpm2_tts.py
+1 -1
View File
@@ -115,7 +115,7 @@ steps:
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export CACHE_DIT_VERSION=1.3.0
- export CACHE_DIT_VERSION=1.5.0
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
+9
View File
@@ -17,6 +17,15 @@ steps:
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
commands:
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'core_model and npu and A3 and omni' --run-level 'core_model'
- label: "Omni · MiniCPM-o 4.5 Perf Test"
timeout_in_minutes: 60
mirror_hardwares: a3_npu_2
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and npu and A3 and omni'
- group: ":card_index_dividers: TTS Model Test"
key: ready-tts-group
+23 -1
View File
@@ -40,9 +40,11 @@ steps:
commands:
- "bash .buildkite/release/scripts/publish-release-wheels-pypi.sh"
# Auto-skip when NIGHTLY=1 (scheduled); otherwise require manual unblock.
- block: "Unblock to build release Docker images"
depends_on: ~
key: block-build-release-images
if: build.env("NIGHTLY") != "1"
- group: "Build release Docker images"
key: "build-release-images"
@@ -63,7 +65,7 @@ steps:
--progress plain \
-f docker/Dockerfile.cuda .
- "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
# re-tag to default image tag and push, just in case arm64 build fails
# Fallback default tag (no arch suffix) if arm64 build fails.
- "docker tag public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT"
- "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT"
env:
@@ -86,6 +88,26 @@ steps:
env:
DOCKER_BUILDKIT: "1"
# NIGHTLY=1 only: promote ECR images to DockerHub nightly tags, then prune.
- label: "Publish nightly multi-arch image to DockerHub"
depends_on:
- build-release-image-x86
- build-release-image-arm64
if: build.env("NIGHTLY") == "1"
key: publish-nightly-images-dockerhub
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/release/scripts/publish-release-images.sh nightly"
- "bash .buildkite/release/scripts/cleanup-nightly-builds.sh"
plugins:
- docker-login#v3.0.0:
username: vllmbot
password-env: DOCKERHUB_TOKEN
env:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
- block: "Publish release images to DockerHub"
key: block-publish-release-images
depends_on:
@@ -0,0 +1,147 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
#
# Clean up old nightly builds from DockerHub (vllm/vllm-omni), keeping the
# newest 14 tags with the nightly- prefix.
set -ex
TAG_PREFIX="nightly-"
REPO="vllm/vllm-omni"
REPO_API_URL="https://hub.docker.com/v2/repositories/${REPO}/tags"
echo "Cleaning up tags with prefix: $TAG_PREFIX in repository: $REPO"
if [ -z "$DOCKERHUB_TOKEN" ]; then
echo "Error: DOCKERHUB_TOKEN environment variable is not set"
exit 1
fi
if [ -z "$DOCKERHUB_USERNAME" ]; then
echo "Error: DOCKERHUB_USERNAME environment variable is not set"
exit 1
fi
echo "Getting DockerHub bearer token..."
set +x
BEARER_TOKEN=$(curl -s -X POST \
-H "Content-Type: application/json" \
-d "{\"username\": \"$DOCKERHUB_USERNAME\", \"password\": \"$DOCKERHUB_TOKEN\"}" \
"https://hub.docker.com/v2/users/login" | jq -r '.token')
set -x
if [ -z "$BEARER_TOKEN" ] || [ "$BEARER_TOKEN" = "null" ]; then
echo "Error: Failed to get DockerHub bearer token"
exit 1
fi
get_all_tags() {
local page=1
local all_tags=""
# Hub `name` is a partial match (not a strict prefix). Use it to narrow the
# list, then enforce startswith() client-side. Exhaust API pages; do not
# stop merely because a page has no matching tags after filtering.
local name_query
name_query=$(printf '%s' "$TAG_PREFIX" | jq -sRr @uri)
while true; do
set +x
local response
response=$(curl -s -H "Authorization: Bearer $BEARER_TOKEN" \
"$REPO_API_URL?page=$page&page_size=100&name=${name_query}")
set -x
local page_count
page_count=$(echo "$response" | jq -r '.results | length')
if [ -z "$page_count" ] || [ "$page_count" = "0" ] || [ "$page_count" = "null" ]; then
break
fi
local tags
tags=$(echo "$response" | jq -r --arg prefix "$TAG_PREFIX" \
'.results[] | select(.name | startswith($prefix)) | "\(.last_updated)|\(.name)"')
if [ -n "$tags" ]; then
all_tags="$all_tags$tags"$'\n'
fi
page=$((page + 1))
done
echo "$all_tags" | sort -r | cut -d'|' -f2
}
delete_tag() {
local tag_name="$1"
echo "Deleting tag: $tag_name"
local delete_url="https://hub.docker.com/v2/repositories/${REPO}/tags/$tag_name"
set +x
local response http_code body
response=$(curl -s -w "\n%{http_code}" -X DELETE \
-H "Authorization: Bearer $BEARER_TOKEN" "$delete_url")
set -x
http_code=$(printf '%s\n' "$response" | tail -n1)
body=$(printf '%s\n' "$response" | sed '$d')
if [ "$http_code" -ge 200 ] && [ "$http_code" -lt 300 ]; then
echo "Successfully deleted tag: $tag_name (HTTP ${http_code})"
return 0
fi
local detail=""
if echo "$body" | jq -e '.detail' > /dev/null 2>&1; then
detail=$(echo "$body" | jq -r '.detail')
elif [ -n "$body" ]; then
detail="$body"
else
detail="(empty response body)"
fi
echo "Error: Failed to delete tag $tag_name (HTTP ${http_code}): ${detail}"
return 1
}
echo "Fetching all tags from DockerHub..."
all_tags=$(get_all_tags)
if [ -z "$all_tags" ]; then
echo "No tags found to clean up"
exit 0
fi
total_tags=$(echo "$all_tags" | wc -l)
echo "Found $total_tags tags"
tags_to_keep=14
tags_to_delete=$((total_tags - tags_to_keep))
if [ $tags_to_delete -le 0 ]; then
echo "No tags need to be deleted (only $total_tags tags found, keeping $tags_to_keep)"
exit 0
fi
echo "Will delete $tags_to_delete old tags, keeping the newest $tags_to_keep"
tags_to_delete_list=$(echo "$all_tags" | tail -n +$((tags_to_keep + 1)))
if [ -z "$tags_to_delete_list" ]; then
echo "No tags to delete"
exit 0
fi
echo "Deleting old tags..."
delete_failures=0
while IFS= read -r tag; do
if [ -n "$tag" ]; then
if ! delete_tag "$tag"; then
delete_failures=$((delete_failures + 1))
fi
sleep 1
fi
done <<< "$tags_to_delete_list"
if [ "$delete_failures" -gt 0 ]; then
echo "Cleanup finished with ${delete_failures} deletion failure(s)"
exit 1
fi
echo "Cleanup completed successfully"
@@ -1,50 +1,87 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
#
# Publish release Docker images from ECR to DockerHub.
# Pulls per-arch images, tags with `latest` and versioned tags, pushes them,
# then creates and pushes a multi-arch manifest.
# Publish Docker images from ECR staging to DockerHub.
#
# Usage:
# publish-release-images.sh # release mode (default)
# publish-release-images.sh release
# publish-release-images.sh nightly
#
# Release tags: latest, v<version>
# Nightly tags: nightly, nightly-<commit>
set -euo pipefail
RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//')
if [ -z "${RELEASE_VERSION}" ]; then
echo "ERROR: release-version metadata not set"
exit 1
fi
MODE="${1:-release}"
COMMIT="$BUILDKITE_COMMIT"
DOCKERHUB_REPO="vllm/vllm-omni"
ECR_REPO="public.ecr.aws/q9t5s3a7/vllm-omni-release-repo"
COMMIT="$BUILDKITE_COMMIT"
PRIMARY_TAGS=()
EXTRA_MANIFEST_TAGS=()
case "${MODE}" in
release)
RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//')
if [ -z "${RELEASE_VERSION}" ]; then
echo "ERROR: release-version metadata not set"
exit 1
fi
PRIMARY_TAGS=("latest" "v${RELEASE_VERSION}")
;;
nightly)
# Commit-pinned manifest reuses nightly-{x86_64,aarch64} arch images.
PRIMARY_TAGS=("nightly")
EXTRA_MANIFEST_TAGS=("nightly-${COMMIT}")
;;
*)
echo "Usage: $0 {release|nightly}"
exit 2
;;
esac
echo "========================================"
echo "Publishing release images v${RELEASE_VERSION}"
echo "Publishing ${MODE} images to ${DOCKERHUB_REPO}"
echo " Commit: ${COMMIT}"
echo " Primary tags: ${PRIMARY_TAGS[*]}"
if [ ${#EXTRA_MANIFEST_TAGS[@]} -gt 0 ]; then
echo " Extra manifests: ${EXTRA_MANIFEST_TAGS[*]}"
fi
echo "========================================"
# Login to ECR to pull staging images
aws ecr-public get-login-password --region us-east-1 | \
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
docker pull "${ECR_REPO}:${COMMIT}-x86_64"
docker pull "${ECR_REPO}:${COMMIT}-aarch64"
SRC_X86="${ECR_REPO}:${COMMIT}-x86_64"
SRC_ARM="${ECR_REPO}:${COMMIT}-aarch64"
docker tag "${ECR_REPO}:${COMMIT}-x86_64" "${DOCKERHUB_REPO}:latest-x86_64"
docker tag "${ECR_REPO}:${COMMIT}-x86_64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64"
docker push "${DOCKERHUB_REPO}:latest-x86_64"
docker push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64"
docker pull "${SRC_X86}"
docker pull "${SRC_ARM}"
docker tag "${ECR_REPO}:${COMMIT}-aarch64" "${DOCKERHUB_REPO}:latest-aarch64"
docker tag "${ECR_REPO}:${COMMIT}-aarch64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
docker push "${DOCKERHUB_REPO}:latest-aarch64"
docker push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
publish_manifest() {
local tag="$1"
local arch_tag_base="$2"
docker manifest rm "${DOCKERHUB_REPO}:${tag}" || true
docker manifest create \
"${DOCKERHUB_REPO}:${tag}" \
"${DOCKERHUB_REPO}:${arch_tag_base}-x86_64" \
"${DOCKERHUB_REPO}:${arch_tag_base}-aarch64"
docker manifest push "${DOCKERHUB_REPO}:${tag}"
}
docker manifest rm "${DOCKERHUB_REPO}:latest" || true
docker manifest rm "${DOCKERHUB_REPO}:v${RELEASE_VERSION}" || true
docker manifest create "${DOCKERHUB_REPO}:latest" "${DOCKERHUB_REPO}:latest-x86_64" "${DOCKERHUB_REPO}:latest-aarch64"
docker manifest create "${DOCKERHUB_REPO}:v${RELEASE_VERSION}" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
docker manifest push "${DOCKERHUB_REPO}:latest"
docker manifest push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}"
for tag in "${PRIMARY_TAGS[@]}"; do
docker tag "${SRC_X86}" "${DOCKERHUB_REPO}:${tag}-x86_64"
docker tag "${SRC_ARM}" "${DOCKERHUB_REPO}:${tag}-aarch64"
docker push "${DOCKERHUB_REPO}:${tag}-x86_64"
docker push "${DOCKERHUB_REPO}:${tag}-aarch64"
publish_manifest "${tag}" "${tag}"
done
ARCH_BASE="${PRIMARY_TAGS[0]}"
for tag in "${EXTRA_MANIFEST_TAGS[@]}"; do
publish_manifest "${tag}" "${ARCH_BASE}"
done
echo ""
echo "Successfully published release images for v${RELEASE_VERSION}"
echo "Successfully published ${MODE} images"
+4 -1
View File
@@ -1,6 +1,6 @@
---
name: precheck-pr
description: Self-check your branch before creating a PR — catch dead code, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub.
description: Self-check your branch before creating a PR — catch dead code, prevent new model-specific Python examples, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub.
---
# PR Pre-Check
@@ -65,6 +65,8 @@ Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type
**Also run the Code-Quality sweep on every PR**, regardless of type or mode: the five diff-scoped checks in [references/code-quality.md](references/code-quality.md) — kwargs fragility, broad-except swallow, `Any`/wrong type hints, hot-path `.clone()`/`deepcopy`, and event-loop blocking — plus the advisory conventions (log level, structured logging, synchronization, cleanup, dependencies, naming) in [checklists.md](references/checklists.md). These count only lines the PR *adds* — the pre-existing backlog across the repo is out of scope.
**Also run the Examples-Policy check on every PR** using [references/examples-policy.md](references/examples-policy.md). Inspect only Python paths introduced by the diff. Treat a new model-, checkpoint-, vendor-, or family-specific Python example as blocking; do not report pre-existing example debt.
### Step 5: Print Report
```
@@ -77,6 +79,7 @@ Pre-check report for <branch>
───────────────── ──────
PR title format ✓
Code quality ⚠ 1 broad except, 2 Any hints
Examples policy ✗ new model-specific Python example
PR desc integrity ✓
Registry/config ✓
Dead code ⚠ 2 warnings
@@ -47,6 +47,19 @@ Diff-scoped sweep for five fragility patterns that are pervasive in the existing
Roll the ⚠/✗ counts into the report as a single **Code quality** dimension row.
---
## All PRs: Examples Policy (run on every PR, quick & full)
Inspect Python files introduced under `examples/` using
[examples-policy.md](examples-policy.md). This is a diff-scoped ratchet: existing
model-specific examples may be edited or removed without triggering the check.
- [ ] **No new model-specific Python examples:** the diff adds, copies, or renames no Python file whose path or behavior is specific to one model, checkpoint, vendor, or model family. ✗ for any violation. Route inference through a shared task/protocol entrypoint, model request behavior through `vllm_omni/model_extras`, and commands through task docs or `recipes/`.
- [ ] **New generic example code is justified:** any new task-neutral Python file under `examples/` is an actual user-facing entrypoint rather than reusable production logic, a benchmark, an app, a setup tool, or a reproducer. ⚠ when the placement is questionable.
Roll the result into the report as a separate **Examples policy** dimension.
### Conventions (eyeball — no grep)
Judgment calls that don't grep cleanly — apply them while reading the diff, not as a mechanical check:
@@ -0,0 +1,68 @@
# Examples Python Policy
Prevent new model-specific Python examples without forcing cleanup of the
existing backlog. Run this check on every PR in quick and full modes.
## Scope
Set the merge base as in the main workflow, then list Python paths introduced
by the PR, including additions, copies, and renames:
```bash
BASE="$(git merge-base HEAD origin/main)"
git diff --name-status --diff-filter=ACR --find-renames --find-copies \
"${BASE}"...HEAD -- 'examples/**/*.py'
```
Inspect destination paths for copies and renames. Existing model-specific files
that are only modified or deleted are grandfathered and must not be reported.
## Blocking classification
Mark ✗ when a newly introduced Python file does any of the following:
- Names or scopes a path to one model, checkpoint, vendor, or model family.
- Hard-codes a model identifier, model-specific prompt, request shape, output
conversion, or launch configuration behind an otherwise generic filename.
- Duplicates an existing task/protocol runner to demonstrate one model.
- Adds a per-model helper below a generic task hub such as `text_to_speech/`.
Do not infer safety from a generic filename alone. For example, a file named
`text_to_audio.py` is still model-specific if it only implements one model's
contract.
## Acceptable additions
A new Python example may pass when it is a genuinely model-neutral,
user-facing task or protocol entrypoint. It must accept model selection through
configuration and keep model-specific behavior out of the script.
Use ⚠ when the file is model-neutral but reusable logic appears misplaced under
`examples/`. Route code by purpose:
| Code purpose | Preferred location |
|---|---|
| Model prompt, defaults, request/output adaptation | `vllm_omni/model_extras/` or production model modules |
| Runnable model command and validation evidence | Task documentation or `recipes/` |
| UI or interactive application | `apps/` |
| Benchmark or evaluator | `benchmarks/` |
| Regression reproducer | `tests/` |
| Download, conversion, or setup utility | `tools/` |
If a model cannot use an existing shared runner, report the missing generic
capability. Do not solve the gap by adding a model-specific Python example.
## Report format
Add one row to the pre-check report:
```text
Examples policy ✓ no new Python example paths
Examples policy ⚠ generic helper belongs in tools/
Examples policy ✗ examples/offline_inference/new_model/end2end.py is model-specific
Examples policy ✗ examples/offline_inference/x_to_y_model_name.py is model-specific
```
For a blocker, name the path, the model-specific behavior, and the appropriate
shared runner or destination. Do not propose deleting unrelated existing
examples as part of the PR.
+8
View File
@@ -97,6 +97,7 @@ If docs and live code disagree, verify the code/tests and report the drift.
| [test-quality-evaluation.md](references/checks/test-quality-evaluation.md) | Tests change, are absent for risky code, or may not exercise production behavior. |
| [tests-docs-checklist.md](references/checks/tests-docs-checklist.md) | Coverage, CI markers, examples, user docs, or PR evidence need review. |
| [verification.md](references/checks/verification.md) | Hardware, a server, or a runnable affected path is available for active verification. |
| [examples-policy.md](../precheck-pr/references/examples-policy.md) | The PR adds, copies, or renames Python under `examples/`; apply the canonical policy shared with `precheck-pr`. |
### Delivery and reviewer coordination
@@ -152,6 +153,13 @@ contract they protect or use only the applicable evidence checks.
Apply every category in [general-checks.md](references/process/general-checks.md) before
lower-priority comments.
If the diff census contains an added, copied, or renamed Python path under
`examples/`, read and apply the canonical
[examples policy](../precheck-pr/references/examples-policy.md). Treat a new
model-specific Python example as blocking. Do not flag model-specific example
debt that the PR only modifies or removes, and do not run the rest of the
author-oriented `precheck-pr` workflow.
For each changed value or behavior, trace:
```text
+7 -2
View File
@@ -11,10 +11,15 @@ build:
post_checkout:
- git fetch --unshallow || true
pre_install:
# `pre_install` runs before RTD upgrades pip, and the seeded pip 23.1
# rejects the `typing_extensions` wheel torch needs (fixed in pip 23.2).
- python -m pip install --upgrade pip
# Pre-install CPU torch so the subsequent `pip install .[docs]` does not
# pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's
# 1 GB memory cap and time out the mkdocs build.
- python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
# 1 GB memory cap and time out the mkdocs build. Wheel-only: `--index-url`
# replaces PyPI, so an sdist fallback cannot reach its build deps. The
# quotes keep YAML from reading `:all:` as a mapping key.
- "python -m pip install --only-binary=:all: --index-url https://download.pytorch.org/whl/cpu torch"
mkdocs:
configuration: mkdocs.yml
+10 -11
View File
@@ -158,7 +158,6 @@ def _build_omni_kwargs(args, quantization=None):
def _generate_image(omni, args, prompt, seed):
"""Generate a single image and return (PIL.Image, time_seconds, memory_gib)."""
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
@@ -176,7 +175,7 @@ def _generate_image(omni, args, prompt, seed):
elapsed = time.perf_counter() - start
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
req_out = OmniRequestOutput.unwrap_result(outputs)
req_out = outputs[0] if isinstance(outputs, list) else outputs
if not req_out.images:
raise ValueError("Could not extract image output from result.")
img = req_out.images[0]
@@ -206,15 +205,15 @@ def _generate_video(omni, args, prompt, seed):
elapsed = time.perf_counter() - start
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
first = outputs[0]
if hasattr(first, "request_output") and isinstance(first.request_output, list):
inner = first.request_output[0]
if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
frames = inner.images[0] if inner.images else None
else:
frames = inner
elif hasattr(first, "images") and first.images:
frames = first.images[0]
if isinstance(outputs, list) and isinstance(outputs[0], OmniRequestOutput):
first = outputs[0]
elif isinstance(outputs, OmniRequestOutput):
first = outputs
else:
raise ValueError("Could not extract video frames from output.")
if hasattr(first, "images"):
frames = first.images[0] if first.images else None
else:
raise ValueError("Could not extract video frames from output.")
+1 -1
View File
@@ -300,7 +300,7 @@ def benchmark(args: argparse.Namespace) -> None:
output_idx = 0
for stage_outputs in omni.generate(all_prompts, sampling_params_list, py_generator=True):
if stage_outputs.final_output_type == "image":
request_output = stage_outputs.request_output
request_output = stage_outputs
request_id = getattr(request_output, "request_id", "")
images = getattr(request_output, "images", [])
@@ -167,9 +167,11 @@ def main() -> None:
]
native_cmd = [
sys.executable,
"examples/offline_inference/text_to_video/text_to_video_lingbot.py",
"examples/offline_inference/text_to_video/text_to_video.py",
"--model",
args.model,
"--model-class-name",
"LingBotVideoPipeline",
"--prompt",
args.prompt,
"--output",
+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
+4 -3
View File
@@ -71,7 +71,7 @@ nav:
- Hybrid Sharded Data Parallel: user_guide/diffusion/parallelism/hsdp.md
- 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
- VAE Parallel: user_guide/diffusion/parallelism/vae_parallel.md
- Attention Backends: user_guide/diffusion/attention_backends.md
- Regional Compilation: user_guide/diffusion/regional_compilation.md
- Frame Interpolation: user_guide/diffusion/frame_interpolation.md
@@ -123,7 +123,7 @@ nav:
- Pipeline Parallel: design/feature/pipeline_parallel.md
- Sequence Parallel: design/feature/sequence_parallel.md
- Tensor Parallel: design/feature/tensor_parallel.md
- VAE Patch Parallelism: design/feature/vae_parallel.md
- VAE Parallel: design/feature/vae_parallel.md
- Attention Optimization:
- Skip-Softmax: design/feature/skip_softmax.md
- Cache-DiT: design/feature/cache_dit.md
@@ -163,7 +163,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
+1 -1
View File
@@ -398,7 +398,7 @@ the source for configuration, compatibility, and implementation details.
* **Composable parallelism:** Diffusion stages can combine [CFG-Parallel](feature/cfg_parallel.md),
[Expert Parallel](feature/expert_parallel.md), [HSDP](feature/hsdp.md),
[Pipeline Parallel](feature/pipeline_parallel.md), [Sequence Parallel](feature/sequence_parallel.md),
[Tensor Parallel](feature/tensor_parallel.md), and [VAE Patch Parallelism](feature/vae_parallel.md)
[Tensor Parallel](feature/tensor_parallel.md), and [VAE Parallel](feature/vae_parallel.md)
according to the pipeline and hardware topology.
* **Attention and cache acceleration:** [Skip-Softmax](feature/skip_softmax.md),
[Cache-DiT](feature/cache_dit.md), and [TeaCache](feature/teacache.md)
+143 -38
View File
@@ -1,42 +1,57 @@
# VAE Patch Parallelism
# VAE Parallel
This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
We use **Qwen-Image** as the reference implementation for decode parallel, and **Wan2.2** for encode parallel.
This document defines the VAE parallelism strategy contract and describes how to
add a strategy to a diffusion model. Tile/patch parallelism is the portable
baseline. Spatial-shard parallelism is a model-specific decode strategy that can
be selected explicitly or through a request-aware policy.
We use **Qwen-Image** as the reference implementation for tile-parallel decode,
**Wan2.2** for tile-parallel encode, and **Wan2.2** as the first spatial-shard
decode implementation. Other VAE families can add spatial strategies later
without changing the public mode meanings defined here.
---
## Table of Contents
- [Overview](#overview)
- [Step-by-Step Implementation (Decode)](#step-by-step-implementation-decode)
- [Encode Parallel Implementation](#encode-parallel-implementation)
- [Strategy Selection Contract](#strategy-selection-contract)
- [Tile/Patch Implementation (Decode)](#tilepatch-implementation-decode)
- [Tile/Patch Encode Implementation](#tilepatch-encode-implementation)
- [Testing](#testing)
- [Reference Implementations](#reference-implementations)
- [Spatially-Sharded Decode](#spatially-sharded-decode)
- [Adding Another VAE Strategy](#adding-another-vae-strategy)
- [Summary](#summary)
---
## Overview
### What is Vae Patch parallel?
### What is VAE parallelism?
**VAE Patch Parallelism** is an acceleration technique for both **encoding** and **decoding**. Instead of processing the entire tensor at once, the tensor is:
VAE parallelism distributes VAE **encoding** or **decoding** across multiple
ranks. The implementation currently has two strategy families:
+ Split into multiple spatial tiles
| Strategy | Unit of work | Communication | Current scope |
|----------|--------------|---------------|---------------|
| Tile/patch | Independent overlapping tiles | Gather, stitch, and optionally broadcast | Encode and decode across supported distributed VAEs |
| Spatial shard | One feature map sharded along a spatial axis | Per-layer halo exchange followed by output gather | Wan decode |
+ Distributed across multiple ranks
Both strategies preserve the model-specific receptive field and reconstruct the
same logical output as single-rank execution. They differ in how work is divided:
tile parallelism assigns complete overlapping tiles to ranks, while spatial
sharding keeps one logical feature map partitioned throughout the decoder.
+ Encoded/Decoded in parallel
This approach can:
+ Merged to reconstruct the final output
+ Distribute computation across multiple devices
+ Reduce peak memory usage per device
+ Accelerate encoding or decoding latency
This approach:
+ Distributes computation across multiple devices
+ Reduces peak memory usage per device
+ Accelerates encoding/decoding latency
The actual benefit depends on shape, topology, communication cost, and VAE
architecture. A strategy must not become the global default without parity and
performance evidence on its supported models.
### When to Use Encode vs Decode Parallel
@@ -45,8 +60,16 @@ This approach:
| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |
### Architecture
We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/decoding.
### Architecture and ownership
`DiffusionParallelConfig.vae_parallel_mode` carries the public strategy request.
The shared registry forwards that value to the VAE, but each distributed VAE
adapter owns its supported modes, request-shape selection, topology validation,
and fallback behavior. Model-specific policy must not be added to the shared
runner or registry.
**DistributedVaeExecutor** is the model-agnostic core for tile-parallel VAE
encoding and decoding.
The executor is model-agnostic and accepts three function parameters:
@@ -76,6 +99,11 @@ This design separates:
+ Model-specific tiling and merging logic
Spatial-shard implementations use the same configured VAE process group but do
not use the tile task executor. They keep the feature map sharded across decoder
layers and own their halo-exchange and final-gather contracts inside the model's
distributed autoencoder implementation.
#### Why split / exec / merge is necessary?
The latent tensor cannot be arbitrarily partitioned.
@@ -92,7 +120,34 @@ Therefore:
+ Merge must perform blending to avoid seams
## Step-by-Step Implementation (Decode)
## Strategy Selection Contract
`vae_parallel_mode="tile"` is the compatibility-preserving global default.
Other values are stable strategy requests, but support remains model-specific.
| Mode | Meaning | Current implementation | Unsupported or ineligible request |
|------|---------|------------------------|-----------------------------------|
| `tile` | Use the existing tile/patch executor | All distributed VAE adapters | Not applicable |
| `auto` | Let the VAE adapter select a validated strategy for each request | Wan decode selects height or width spatial sharding | Wan uses tile when ineligible; unsupported VAE families reject or document their fallback |
| `spatial_shard_height` | Force spatial-shard decode along height | Wan decode | Wan warns and uses tile when ineligible; unsupported VAE families reject |
| `spatial_shard_width` | Force spatial-shard decode along width | Wan decode | Wan warns and uses tile when ineligible; unsupported VAE families reject |
The shared configuration layer validates the mode name. It does not decide
whether a model supports that mode. A VAE adapter adding a mode must:
1. Keep `tile` as an exact feature-off path.
2. Select a strategy per request rather than from stale process-global state.
3. Validate input rank/shape and process-group membership before collectives.
4. Reject or explicitly fall back from unsupported combinations; never silently
reinterpret one forced strategy as another.
5. Preserve direct encode/decode behavior when Diffusers does not select a tiled
path.
`auto` is a policy hook, not a universal heuristic. Each VAE family may select
from only the strategies it has validated. Adding another family must not change
Wan's selector or the meanings of the explicit modes.
## Tile/Patch Implementation (Decode)
### Step 1: Implement DistributedAutoencoderKLQwenImage
`QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
@@ -251,7 +306,7 @@ class YourModelPipeline(nn.Module):
+ ).to(self.device)
```
## Encode Parallel Implementation
## Tile/Patch Encode Implementation
For models that require VAE encoding (e.g., Image-to-Video), you can also parallelize the encode operation. We use **Wan2.2** as the reference implementation.
@@ -449,18 +504,24 @@ Complete examples in the codebase:
---
## Spatially-Sharded Decode (Wan)
## Spatially-Sharded Decode
The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE additionally supports a **spatially-sharded decode** backend, selected through `DiffusionParallelConfig.vae_parallel_mode` (`"spatial_shard_height"` or `"spatial_shard_width"`; default is `"tile"`).
The tile-parallel executor above assigns independent spatial **tiles** to ranks.
A spatial-shard backend instead partitions one global decoder feature map along
height or width. Wan is the first supported implementation.
`DiffusionParallelConfig.vae_parallel_mode="tile"` remains the default. For Wan,
opt into request-shape selection with `"auto"`, or force an axis with
`"spatial_shard_height"` or `"spatial_shard_width"`.
### How it differs from tile parallel
| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
|--------|--------------------------|-----------------------------------------------|
| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W |
| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) |
| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
| Scope | Decode + encode | Decode only |
| Aspect | Tile parallel (`"tile"`) | Automatic (`"auto"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
|--------|--------------------------|----------------------|-----------------------------------------------|
| Unit of work | Independent overlapping tiles | Chooses tile or a spatial shard per request | A single global feature map sharded along H or W |
| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Depends on the selected path | Per-conv **halo exchange** of boundary rows/cols (P2P) |
| Output assembly | Blend overlapping tiles | Depends on the selected path | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
| Scope | Decode + encode | Wan decode selection; encode remains tiled | Wan decode only |
Spatial-shard decode swaps the decoder's spatial convolutions/padding for halo-exchanging variants (`WanDistConv2d`, `WanDistCausalConv3d`, `WanDistZeroPad2d`) so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Implementation lives in `vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py`.
@@ -476,21 +537,65 @@ serve.py (--vae-parallel-mode) / OmniEngineArgs
-> DistributedAutoencoderKLWan.tiled_decode dispatch
```
`DistributedAutoencoderKLWan._spatial_shard_decode_enabled()` gates the path: it requires distributed decode to be enabled, a 5D latent, and `vae_patch_parallel_size == DiT group size`. Otherwise it logs a warning and falls back to tile-parallel decode.
`DistributedAutoencoderKLWan._spatial_shard_decode_enabled()` gates the path: it
requires distributed decode to be enabled, a 5D latent, and
`vae_patch_parallel_size == DiT group size`. The method is reached only after
Diffusers has selected tiled decode for the request. In `auto` mode the longer
axis that can cover every rank is selected, with width winning ties. If neither
axis can cover the group, auto mode uses tile decode. This policy is cached only
through normal kernel shape caches; it is reevaluated for every request.
### Notes
- The decoder is patched **in place** the first time spatial-shard decode runs and is bound to a single split dimension for the lifetime of the VAE instance.
- The decoder is patched **in place** the first time spatial-shard decode runs. Its wrappers consult a context-local split dimension and retain an exact direct fallback, allowing later requests to select tile, height, or width safely.
- Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `distributed_cuda`).
## Adding Another VAE Strategy
Add future strategies through the distributed autoencoder that owns the model
semantics. Do not add model-name dispatch to the shared runner, registry, or
`DistributedVaeExecutor`.
For a new VAE family:
1. **Declare capability**: document which modes, operations, shapes, dtypes, and
process-group layouts the adapter supports.
2. **Implement model-local selection**: keep `auto` policy and explicit-mode
dispatch in the distributed autoencoder. The selector may use request shape,
but must be deterministic and side-effect free.
3. **Preserve the baseline**: retain exact tile/direct fallbacks and state-dict
compatibility. Structural wrappers must not register duplicate parameters or
alter checkpoint names.
4. **Define communication**: derive shards from the configured VAE group, keep
collectives symmetric, specify padding/halo ownership, and return the same
output placement as the tile path.
5. **Bound request state**: make the selected axis and scratch buffers
request-scoped, release temporary buffers on success and failure, and support
alternating modes in one long-lived worker.
6. **Validate independently**: compare each explicit strategy against the
single-rank reference, then test `auto`, tile fallback, unsupported topology,
and a mixed request sequence. Record per-rank peak memory and decode latency
before documenting a preferred mode.
Promote a selector into shared infrastructure only after at least two VAE
families use the same inputs, invariants, fallback behavior, and lifecycle. Until
then, shared code owns transport and configuration while each model adapter owns
policy.
---
## Summary
Adding VAE Patch Parallel support to diffusion model:
Adding VAE parallel support to a diffusion model:
1. **Implement Distributed VAE** - Inherit from base VAE class and `DistributedVaeMixin`
2. **Decode Parallel** - Refactor `tiled_decode` into `tile_split`/`tile_exec`/`tile_merge`
3. **Encode Parallel** (optional) - Implement `encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models
4. **Change VAE model in pipeline** - Use the distributed version
5. **Test** - Verify numerical consistency with `vae_patch_parallel_size=1` vs `N`
1. **Implement the distributed VAE** - Inherit from the base VAE class and
`DistributedVaeMixin`.
2. **Keep tile as the baseline** - For tile decode, refactor `tiled_decode` into
`tile_split`/`tile_exec`/`tile_merge`.
3. **Add encode parallelism when needed** - Implement
`encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models.
4. **Add model-specific strategies locally** - Own selectors, topology checks,
communication, and fallback in the distributed autoencoder.
5. **Wire the pipeline** - Replace the pipeline VAE with its distributed version.
6. **Test strategy and lifecycle parity** - Compare one rank with every supported
strategy, including feature-off, fallback, and mixed-request sequences.
+1 -1
View File
@@ -48,7 +48,7 @@ implementation contract; it is not, by itself, a general support claim.
- [Pipeline Parallel](feature/pipeline_parallel.md)
- [Sequence Parallel](feature/sequence_parallel.md)
- [Tensor Parallel](feature/tensor_parallel.md)
- [VAE Patch Parallelism](feature/vae_parallel.md)
- [VAE Parallel](feature/vae_parallel.md)
#### Attention optimization
+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
+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.
@@ -10,7 +10,7 @@ This guide covers the parallelism methods in vLLM-Omni for speeding up diffusion
| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation |
| **[Pipeline Parallelism](pipeline_parallel.md)** | Splits the denoising transformer block-wise across sequential GPU stages to reduce per-GPU model memory |
| **[VAE Parallelism](vae_parallelism.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
| **[VAE Parallel](vae_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs |
| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) |
@@ -1,4 +1,4 @@
# VAE Parallelism Guide
# VAE Parallel Guide
## Table of Content
@@ -113,7 +113,7 @@ In `DiffusionParallelConfig`:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. |
| `vae_parallel_mode` | str | `"tile"` | VAE parallel decode strategy: `"tile"` (default tile/patch parallel decode), `"spatial_shard_height"`, or `"spatial_shard_width"` (spatially-sharded decode, Wan only). See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). |
| `vae_parallel_mode` | str | `"tile"` | VAE parallel decode strategy: `"tile"`, `"auto"` (Wan request-shape selection), `"spatial_shard_height"`, or `"spatial_shard_width"`. Spatial sharding is currently Wan-only. See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). |
Additional requirements:
@@ -128,7 +128,7 @@ Additional requirements:
## Spatially-Sharded Decode (Wan)
The default `vae_parallel_mode="tile"` distributes whole tiles across ranks. For the **Wan** VAE there is an alternative decode strategy, **spatially-sharded decode**, selected via `vae_parallel_mode="spatial_shard_height"` or `vae_parallel_mode="spatial_shard_width"`.
The default `vae_parallel_mode="tile"` keeps patch/tile parallel decode unchanged. For Wan, set `vae_parallel_mode="auto"` to select spatial-shard decode only after a request is already large enough to enter Diffusers' tiled-decode path and the VAE uses the full DiT decode group. Auto mode splits the longer eligible spatial axis to reduce halo traffic. Direct-decode requests and partial groups fall back to tile decode. Use `"spatial_shard_height"` or `"spatial_shard_width"` to force an axis.
Instead of assigning independent tiles to ranks, spatial-shard decode shards the decoder feature maps along the height (`spatial_shard_height`) or width (`spatial_shard_width`) dimension and exchanges halo rows/columns between neighboring ranks around the spatial convolutions. This keeps the receptive field correct across shard boundaries, so the result matches the single-GPU decode within numerical tolerance.
@@ -141,7 +141,7 @@ omni = Omni(
parallel_config=DiffusionParallelConfig(
tensor_parallel_size=2,
vae_patch_parallel_size=2, # must match the DiT group size
vae_parallel_mode="spatial_shard_width", # or "spatial_shard_height"
vae_parallel_mode="auto", # opt in to Wan request-shape selection
),
)
```
@@ -152,14 +152,14 @@ Or from the CLI / serving entrypoint:
vllm serve Wan-AI/Wan2.1-T2V-1.3B-Diffusers --omni \
--tensor-parallel-size 2 \
--vae-patch-parallel-size 2 \
--vae-parallel-mode spatial_shard_width
--vae-parallel-mode auto
```
**Constraints and behavior:**
- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other models ignore `spatial_shard_*` modes.
- It requires `vae_patch_parallel_size` to **match the DiT process group size**. If it does not, the VAE logs a warning and **falls back to tile-parallel decode** at runtime.
- `spatial_shard_height` and `spatial_shard_width` are mutually exclusive for a given VAE instance (the decoder is patched in place for a single split dimension).
- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other VAE families retain their existing mode handling and may reject explicit Wan spatial modes.
- It requires `vae_patch_parallel_size` to **match the DiT process group size**. `auto` quietly falls back to tile decode for a partial group; an explicitly forced spatial mode logs a warning.
- The decoder's context-aware wrappers retain the direct path, so one process can alternate small/large and portrait/landscape requests safely.
For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark:
+1 -1
View File
@@ -61,7 +61,7 @@ Memory optimization methods help reduce GPU memory usage, enabling inference on
|--------|-------------|----------|
| **[CPU Offload](diffusion/cpu_offload.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
| **[Quantization](quantization/overview.md)** | Reduces transformer stages from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss |
| **[VAE Parallelism](diffusion/parallelism/vae_parallelism.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak |
| **[VAE Parallel](diffusion/parallelism/vae_parallel.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak |
### Extensions
@@ -12,6 +12,7 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
|---|---|---|---|---|---|
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) |
| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
@@ -35,6 +36,25 @@ python text_to_video.py \
--output t2v_out.mp4
```
### LingBot-Video
The shared runner recognizes both official LingBot checkpoint IDs and selects
the native `LingBotVideoPipeline` automatically:
```bash
python text_to_video.py \
--model robbyant/lingbot-video-dense-1.3b \
--prompt "a robotic arm picks up a red block" \
--height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
--guidance-scale 3.0 --flow-shift 3.0 --fps 24 \
--output lingbot_t2v.mp4
```
Use `robbyant/lingbot-video-moe-30b-a3b` for the MoE checkpoint. For local
paths without `lingbot` in the directory name, add
`--model-class-name LingBotVideoPipeline`. Model-only options use
`--extra-body`, for example `'{"batch_cfg": true, "output_type": "np"}'`.
### LTX-2
```bash
+15 -15
View File
@@ -12,20 +12,20 @@ For the internal architecture and backend extension points, see the
| Mode | Guide | Description | Methods |
|------|-------|-------------|---------|
| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8, MXFP4 W4A4 |
| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, BitsAndBytes W4, MXFP8 W8A8, MXFP4 W4A4 |
| Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA |
| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale |
| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale |
## Hardware Support
| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | GGUF | AutoRound | msModelSlim |
|--------|----------|-----------|----------|------------|------------|------|-----------|-------------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ⭕ | ⭕ | ❌ |
| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ |
| Device | FP8 W8A8 | Int8 W8A8 | BitsAndBytes W4 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | AutoRound | msModelSlim |
|--------|----------|-----------|------------------|----------|------------|------------|-----------|-------------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ |
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
| AMD ROCm | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ✅ | ⭕ | ❌ |
| Intel XPU | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
guide. FP8 on Ampere may use a weight-only path where available.
@@ -43,10 +43,10 @@ otherwise.
|--------|-------|------|----------------|--------|
| FP8 W8A8 | [FP8](fp8.md) | Online W8A8 or checkpoint FP8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
| Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
| BitsAndBytes W4 | [BitsAndBytes](bitsandbytes.md) | Online W4 weight-only | Z-Image-Turbo; Qwen-Image/Wan2.2 not validated | Validated for Z-Image-Turbo |
| ModelOpt | [ModelOpt](modelopt.md) | Pre-quantized FP8 checkpoints | Qwen-Image, Z-Image, FLUX.2, HunyuanImage-3.0 | Validated for ModelOpt FP8 diffusion checkpoints |
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Online W8A8 or offline pre-quantized | Wan2.2-T2V-A14B, I2V-A14B, TI2V-5B | Ascend NPU only; validated for Wan2.2 |
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | `mxfp4`: online single-scale only; `mxfp4_dualscale`: online or offline dual-scale (offline recommended) | Wan2.2-T2V-A14B, I2V-A14B | Ascend NPU only; validated for Wan2.2 A14B cascade models; TI2V-5B not supported; offline `mxfp4_dualscale` uses calibrated `mul_scale` for best accuracy |
| GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists |
| AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven |
| msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path |
@@ -61,9 +61,9 @@ in BF16 unless the model guide explicitly adds support.
|--------|-------|-------|----------------|--------|
| ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
| Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| MXFP8 | [MXFP8](mxfp8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| MXFP4 | [MXFP4](mxfp4.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints |
| msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
@@ -76,10 +76,10 @@ attached to the intended stage rather than applied globally.
|--------|-------|-------|----------------|--------|
| FP8 | [FP8](fp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
| Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
| BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Stage-specific transformer or DiT module | BAGEL, GLM-Image | Not validated |
| ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation |
| MXFP8 | [MXFP8](mxfp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
| MXFP4 | [MXFP4](mxfp4.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
| GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed |
| AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed |
| msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation |
@@ -106,7 +106,7 @@ config = build_quant_config({
| Component | Default quantized? | Notes |
|-----------|--------------------|-------|
| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, MXFP8, MXFP4, GGUF, AutoRound, and msModelSlim |
| Diffusion transformer | Yes | Primary target for FP8, Int8, BitsAndBytes, ModelOpt, MXFP8, MXFP4, AutoRound, and msModelSlim |
| Text encoder | No | Keep BF16 unless a method-specific guide documents support |
| VAE | No | Keep BF16; storage-only paths are method-specific |
| Scheduler/tokenizer | No | Loaded from the base model repository |
@@ -139,8 +139,8 @@ from vllm_omni.quantization import build_quant_config
build_quant_config("fp8")
build_quant_config({"method": "fp8", "activation_scheme": "static"})
build_quant_config("bitsandbytes")
build_quant_config("auto-round", bits=4, group_size=128)
build_quant_config({"method": "gguf", "gguf_model": "/path/to/model.gguf"})
build_quant_config({"transformer": {"method": "fp8"}, "vae": None})
build_quant_config(None)
```
@@ -131,7 +131,7 @@ def build_sampling_params() -> list[SamplingParams]:
def save_stage_output(stage_outputs, output_dir: str) -> None:
output = stage_outputs.request_output
output = stage_outputs
request_id = output.request_id
if stage_outputs.final_output_type == "text":
text_output = output.outputs[0].text
@@ -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,
@@ -104,7 +104,7 @@ def main(args):
os.makedirs(output_dir, exist_ok=True)
for stage_outputs in omni_outputs:
output = stage_outputs.request_output
output = stage_outputs
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -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",
@@ -228,7 +228,7 @@ async def main():
if not outputs:
raise ValueError("No output produced from omni.generate()")
first_out = outputs[0].request_output
first_out = outputs[0]
req_out: OmniRequestOutput = first_out
# Verify trajectory data (from custom pipeline)
@@ -741,7 +741,7 @@ def make_mmu_prompt(
def iter_mm_outputs(outputs: list[Any]):
for omni_out in outputs:
req_out = getattr(omni_out, "request_output", None)
req_out = omni_out
req_list = req_out if isinstance(req_out, list) else [req_out]
for item in req_list:
if item is None:
@@ -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}")
@@ -303,7 +290,7 @@ def main():
)
img_idx = 0
for req_output in omni_outputs:
ro = getattr(req_output, "request_output", None)
ro = req_output
stage_id = getattr(req_output, "stage_id", None)
# AR stage text — each CompletionOutput.text is already a delta when
@@ -621,7 +621,7 @@ def main():
images = getattr(output, "images", None)
if images:
break
req_out = getattr(output, "request_output", None)
req_out = output
images = getattr(req_out, "images", None) if req_out is not None else None
if images:
break
@@ -646,8 +646,8 @@ def main():
if frames.multimodal_output and "audio" in frames.multimodal_output:
audio = frames.multimodal_output["audio"]
audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
if frames.is_pipeline_output and frames.request_output is not None:
inner_output = frames.request_output
if frames.is_pipeline_output and frames is not None:
inner_output = frames
if isinstance(inner_output, OmniRequestOutput):
if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
audio = inner_output.multimodal_output["audio"]
@@ -0,0 +1,117 @@
"""
Offline inference example: MammothModa2 image summarization (single AR stage).
Example:
uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
--model path/to/MammothModa2-Preview \
--deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \
--image /path/to/input.jpg \
--question "Please summarize the content of this image."
"""
from __future__ import annotations
import argparse
import os
from PIL import Image
from vllm import SamplingParams
from vllm.multimodal.image import convert_image_mode
from vllm_omni import Omni
DEFAULT_SYSTEM = "You are a helpful assistant."
DEFAULT_QUESTION = "Please summarize the content of this image."
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.")
parser.add_argument("--image", type=str, required=True, help="Path to input image.")
parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
parser.add_argument(
"--max-tokens",
type=int,
default=512,
help="Max new tokens to generate.",
)
parser.add_argument("--temperature", type=float, default=0.2)
parser.add_argument("--top-p", type=float, default=0.9)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--trust-remote-code", action="store_true")
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
return parser.parse_args()
def build_prompt(system: str, question: str) -> str:
return (
f"<|im_start|>system\n{system}<|im_end|>\n"
"<|im_start|>user\n"
"<|vision_start|><|image_pad|><|vision_end|>"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
def main() -> None:
args = parse_args()
if not os.path.exists(args.image):
raise FileNotFoundError(f"Image file not found: {args.image}")
pil_image = Image.open(args.image)
image_data = convert_image_mode(pil_image, "RGB")
prompt = build_prompt(args.system, args.question)
omni = Omni(
model=args.model,
deploy_config=args.deploy_config,
trust_remote_code=args.trust_remote_code,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
)
try:
sp = SamplingParams(
temperature=float(args.temperature),
top_p=float(args.top_p),
top_k=-1,
max_tokens=int(args.max_tokens),
seed=int(args.seed),
detokenize=True,
)
# NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None].
# Consume it inside the try block so the worker isn't closed early.
outputs = list(
omni.generate(
[
{
"prompt": prompt,
"multi_modal_data": {"image": image_data},
"additional_information": {"omni_task": ["chat"]},
}
],
[sp],
)
)
finally:
omni.close()
lines: list[str] = []
for stage_outputs in outputs:
ro = stage_outputs
text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
lines.append("answer:\n")
lines.append(text.strip() + "\n")
lines.append("\n")
print("\n".join(lines))
if __name__ == "__main__":
main()
@@ -297,7 +297,7 @@ def main(args):
os.makedirs(output_dir, exist_ok=True)
for stage_outputs in omni_outputs:
output = stage_outputs.request_output
output = stage_outputs
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -347,7 +347,7 @@ def main(args):
omni.start_profile(stages=args.profiler_stages)
for stage_outputs in omni.generate(prompts, sampling_params_list):
output = stage_outputs.request_output
output = stage_outputs
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -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,
@@ -0,0 +1,229 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline speech-to-speech with NVIDIA-NemotronLabs-VoiceChat-11B.
Feeds a 16 kHz user utterance (plus a spoken-style system prompt) through the
3-stage nemotron_voicechat pipeline and writes the agent's reply as text and a
22.05 kHz WAV.
Usage:
python end2end.py \
--checkpoint /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
--wav /path/to/user_question_16k.wav \
--output-dir results/nemotron_voicechat
The system prompt defaults to the NeMo reference default. The text tokenizer
resolves from the ``nvidia/NVIDIA-Nemotron-Nano-9B-v2`` HF id (or the
``NEMOTRON_VOICECHAT_LLM_PATH`` env var for air-gapped runs).
Sizing the input WAV: the timeline is frame-locked, so the reply budget IS the
input duration — the acoustic channel trails the text channel, and a reply
that has not finished when the frames run out is truncated SILENTLY (no error;
ASR of the WAV will just disagree with the emitted text). Leave generous
trailing silence after the question (a question ending at ~4.5 s truncated in
an 8 s WAV but completed cleanly in a 16 s one). The script warns when the
text channel is still emitting non-PAD tokens near the last frame.
"""
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
import numpy as np
import soundfile as sf
# Warn if the text channel is still speaking within this many final frames
# (12 frames = ~1 s) — the frame-locked reply likely got cut off.
TRUNCATION_GUARD_FRAMES = 12
# The NeMo reference offline script's default system prompt.
DEFAULT_SYSTEM_PROMPT = (
"You are an AI voice assistant developed by NVIDIA. "
"Your name is NVIDIA Voice Chat. "
"Answer in a spoken, conversational style rather than a written one. "
"Do not repeat the same sentence over and over again. "
"Start the conversation by greeting the user."
)
def load_wav_16k_mono(path: str) -> np.ndarray:
wav, sr = sf.read(path, dtype="float32")
if wav.ndim > 1:
wav = wav.mean(axis=1)
if sr != 16000:
from vllm.multimodal.audio import resample_audio_scipy
wav = resample_audio_scipy(wav, orig_sr=sr, target_sr=16000)
return wav.astype(np.float32)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--checkpoint", required=True, help="NVIDIA-NemotronLabs-VoiceChat-11B directory")
parser.add_argument("--wav", required=True, help="user speech (any sr; resampled to 16 kHz mono)")
parser.add_argument("--output-dir", default="results/nemotron_voicechat")
parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT)
parser.add_argument(
"--stage-configs-path",
default=None,
help=(
"Deploy yaml override. Pass vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml "
"for the async-chunk streaming pipeline (stages overlap; audio starts decoding "
"before the thinker finishes). Default: the bundled non-streaming yaml (bit-exact "
"parity path)."
),
)
args = parser.parse_args()
if not Path(args.wav).is_file():
raise FileNotFoundError(f"Input WAV not found: {args.wav}")
ckpt = Path(args.checkpoint)
raw_cfg = json.loads((ckpt / "config.json").read_text())
stt_cfg = raw_cfg["model"]["stt"]["model"]
from transformers import AutoTokenizer
from vllm.sampling_params import SamplingParams
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
compute_acoustic_frame_count,
)
tok_ref = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or stt_cfg.get(
"pretrained_llm", "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
)
tokenizer = AutoTokenizer.from_pretrained(tok_ref, trust_remote_code=False)
bos_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("bos_token", "<s>"))
eos_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("eos_token", "</s>"))
pad_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("pad_token", "<SPECIAL_12>"))
# NeMo encode_system_prompt: [bos] + tokens + [eos].
prompt_ids = [bos_id] + tokenizer.encode(args.system_prompt, add_special_tokens=False) + [eos_id]
wav = load_wav_16k_mono(args.wav)
n_frames = compute_acoustic_frame_count(stt_cfg, int(wav.shape[0]))
print(f"input: {wav.shape[0]} samples @16kHz -> {n_frames} acoustic frames; prompt {len(prompt_ids)} tokens")
inputs = {
# vLLM prompt = system prompt + one placeholder position (acoustic
# frame 0); the thinker replaces every position with fused embeddings.
"prompt_token_ids": prompt_ids + [pad_id],
"additional_information": {
# Plain float list: raw ndarrays do not survive the engine-core
# message serialization.
"nvc_audio": wav.astype(np.float32).tolist(),
"nvc_sr": 16000,
"nvc_prompt_token_ids": prompt_ids,
"nvc_expected_frames": n_frames,
},
}
sampling_params_list = [
# thinker: greedy, frame-locked generation length. NO min_tokens: the
# tokenizer's eos_token is <SPECIAL_12> — the frame-locked PAD/silence
# token — and vLLM's min_tokens masks "EOS" to -inf, which would forbid
# silence and force the agent to babble through the whole utterance.
SamplingParams(temperature=0.0, max_tokens=n_frames, ignore_eos=True, seed=0),
# talker: placeholder loop, stops on the stage's stop token. 16383:
# the 1-token placeholder prompt + max_tokens must fit the stage's
# max_model_len (16384).
SamplingParams(temperature=0.0, max_tokens=16383, stop_token_ids=[1], detokenize=False, seed=0),
# code2wav: single decode step.
SamplingParams(temperature=0.0, max_tokens=1, detokenize=False, seed=0),
]
omni_kwargs: dict = {}
if args.stage_configs_path:
omni_kwargs["stage_configs_path"] = args.stage_configs_path
omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs)
outputs = omni.generate(inputs, sampling_params_list)
out_dir = Path(args.output_dir)
out_dir.mkdir(parents=True, exist_ok=True)
stem = Path(args.wav).stem
agent_text = None
audio = None
sr = 22050
for stage_outputs in outputs:
# One entry per final stage (text from the thinker, audio from code2wav);
# mirrors the qwen2_5_omni offline example's collection loop.
output = getattr(stage_outputs, "request_output", stage_outputs)
final_type = getattr(stage_outputs, "final_output_type", None)
completions = getattr(output, "outputs", None) or []
if not completions:
continue
completion = completions[0]
if final_type in (None, "text"):
text = getattr(completion, "text", None)
if text:
agent_text = text
token_ids = getattr(completion, "token_ids", None)
if token_ids:
(out_dir / f"{stem}_text_tokens.json").write_text(json.dumps(list(map(int, token_ids))))
# Frame lock: the reply budget IS the input length. If the text
# channel is still mid-sentence when the frames run out, the
# spoken answer is silently truncated — warn instead.
tail = list(map(int, token_ids))[-TRUNCATION_GUARD_FRAMES:]
if any(t != pad_id for t in tail):
print(
f"WARNING: the text channel emits non-PAD tokens within the last "
f"{TRUNCATION_GUARD_FRAMES} frames ({TRUNCATION_GUARD_FRAMES * 80} ms) — the reply "
"likely ran out of acoustic frames and is cut off mid-sentence. "
"Pad the input WAV with more trailing silence and rerun."
)
# Best-effort: dump the function-channel debug timeline if the
# engine surfaced the request info on this output.
info = getattr(output, "additional_information", None)
func_timeline = info.get("nvc_function_tokens") if isinstance(info, dict) else None
if func_timeline is not None:
values = func_timeline.tolist() if hasattr(func_timeline, "tolist") else list(func_timeline)
(out_dir / f"{stem}_function_tokens.json").write_text(json.dumps([int(v) for v in values]))
print(f"function-channel timeline: {len(values)} entries")
if final_type in (None, "audio"):
# multimodal_output is a MultimodalPayload (Mapping) — duck-type it.
mm = getattr(completion, "multimodal_output", None)
get = getattr(mm, "get", None)
if not callable(get):
continue
candidate = None
for key in ("audio", "model_outputs", "wav", "waveform"):
value = get(key)
if value is not None:
candidate = value
break
if isinstance(candidate, list):
candidate = candidate[0] if candidate else None
if candidate is not None:
arr = np.asarray(
candidate.float().cpu() if hasattr(candidate, "cpu") else candidate, dtype=np.float32
).reshape(-1)
if arr.size > 0:
audio = arr
sr_val = get("sr")
if sr_val is not None:
sr_item = sr_val[0] if isinstance(sr_val, list) else sr_val
sr = int(sr_item.item() if hasattr(sr_item, "item") else sr_item)
if agent_text is not None:
# The frame-locked text channel emits PAD tokens on silent frames;
# strip them (NeMo's post-inference detokenizer does the same).
pad_token = stt_cfg.get("pad_token", "<SPECIAL_12>")
agent_text = " ".join(agent_text.replace(pad_token, " ").split())
text_path = out_dir / f"{stem}_output.txt"
text_path.write_text(agent_text + "\n")
print(f"agent text -> {text_path}\n{agent_text}")
else:
print("WARNING: no agent text in outputs")
if audio is not None and audio.size > 0:
wav_path = out_dir / f"{stem}_output.wav"
sf.write(wav_path, audio, sr)
print(f"agent audio -> {wav_path} ({audio.size / sr:.2f}s @ {sr} Hz)")
else:
raise RuntimeError("No audio produced by the code2wav stage.")
if __name__ == "__main__":
main()
@@ -392,7 +392,7 @@ def main(args):
total_requests = len(prompts)
processed_count = 0
for stage_outputs in omni_generator:
output = stage_outputs.request_output
output = stage_outputs
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -367,7 +367,7 @@ def main(args):
print(f"query type: {args.query_type}")
for stage_outputs in omni_generator:
output = stage_outputs.request_output
output = stage_outputs
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -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",
@@ -237,7 +237,7 @@ async def run_single_request(
sampling_params_list=sampling_params_list,
output_modalities=output_modalities,
):
output = omni_output.request_output
output = omni_output
if omni_output.final_output_type == "text":
if stage_0_first_output_ts is None:
stage_0_first_output_ts = time.perf_counter()
@@ -26,7 +26,6 @@ import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
@@ -352,8 +351,8 @@ def main():
else:
print("[Profiler] No valid profiling data returned.")
# Extract output from result
output = OmniRequestOutput.unwrap_result(result)
# omni.generate() returns a list for sync calls; unwrap single-result list.
output = result[0] if isinstance(result, list) else result
if not output.images:
raise ValueError("No video frames found in OmniRequestOutput.")
@@ -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)
@@ -345,7 +335,7 @@ def main(args):
# Process outputs from each stage
for stage_idx, stage_outputs in enumerate(omni_outputs):
# request_output may be a single object or a list
raw_output = stage_outputs.request_output
raw_output = stage_outputs
outputs_list = raw_output if isinstance(raw_output, list) else [raw_output]
if stage_outputs.final_output_type == "text":
@@ -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,
@@ -43,6 +43,7 @@ from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
AUDIOX_TASKS = ("t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m")
@@ -388,9 +389,9 @@ def main():
raise ValueError("No output generated from omni.generate()")
output = outputs[0]
if not hasattr(output, "request_output") or not output.request_output:
if not isinstance(output, OmniRequestOutput) or not output:
raise ValueError("No request_output found in OmniRequestOutput")
request_output = output.request_output
request_output = output
if not hasattr(request_output, "multimodal_output"):
raise ValueError("No multimodal_output found in request_output")
@@ -115,7 +115,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
)
images_outputs = []
for output in outputs:
req_out = output.request_output
req_out = output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
images = req_out.images
@@ -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
@@ -656,7 +648,7 @@ def main():
images = getattr(output, "images", None)
if images:
break
req_out = getattr(output, "request_output", None)
req_out = output
images = getattr(req_out, "images", None) if req_out is not None else None
if images:
break
@@ -372,18 +372,10 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
Streaming requires `async_chunk: true` in the stage config.
### Word Timestamps
Generate a WAV offline and a JSON sidecar with word-level timestamps from
`Qwen/Qwen3-ForcedAligner-0.6B`:
```bash
python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--forced-aligner Qwen/Qwen3-ForcedAligner-0.6B \
--text "Hello world." \
--output-dir /tmp/qwen3_tts_timestamps
```
The script writes `qwen3_tts_word_timestamps.wav` and
`qwen3_tts_word_timestamps.json`. On machines without a local CUDA toolkit,
set `VLLM_USE_FLASHINFER_SAMPLER=0` to avoid FlashInfer sampler JIT.
Word-level timestamps are currently a serving-path feature: launch
`vllm-omni serve` with `--forced-aligner` and request `word_timestamps`
(see `examples/online_serving/text_to_speech/README.md`). An offline
example is not provided in this release.
### Batched decoding
The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, point `--stage-configs-path` at a stage configs YAML with the desired values (see `vllm_omni/model_executor/stage_configs/` for templates):
@@ -168,7 +168,7 @@ def run_e2e():
print(f"Received {len(outputs)} outputs.")
for i, output in enumerate(outputs):
try:
ro = output.request_output
ro = output
if ro is None:
print("No request_output found.")
continue
@@ -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
@@ -179,7 +176,7 @@ def main(args):
t_start = time.perf_counter()
for stage_outputs in omni.generate(inputs):
request_output = stage_outputs.request_output
request_output = stage_outputs
if request_output is None or not request_output.outputs:
continue
_save_wav(
@@ -215,7 +212,7 @@ async def main_streaming(args):
sample_rate = None
async for stage_output in omni.generate(prompt, request_id=request_id):
mm = stage_output.request_output.outputs[0].multimodal_output
mm = stage_output.outputs[0].multimodal_output
if not stage_output.finished:
t_now = time.perf_counter()
sample_rate, n = _collect_audio_chunk(mm, all_audio_chunks, sample_rate)
@@ -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,
)
@@ -137,7 +137,7 @@ def main(args) -> None:
params_list = sampling_params
for stage_outputs in omni.generate(inputs, params_list):
for i, req_output in enumerate(stage_outputs.request_output):
for i, req_output in enumerate(stage_outputs):
for j, out in enumerate(req_output.outputs):
mm = out.multimodal_output
if mm is None:
@@ -135,7 +135,7 @@ def run_e2e():
print(f"Received {len(outputs)} outputs.")
for i, output in enumerate(outputs):
try:
ro = output.request_output
ro = output
if ro is None:
print("No request_output found.")
continue
@@ -391,7 +391,7 @@ def main(args):
for batch_start in range(0, len(inputs), batch_size):
batch = inputs[batch_start : batch_start + batch_size]
for stage_outputs in omni.generate(batch):
output = stage_outputs.request_output
output = stage_outputs
_save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
@@ -412,7 +412,7 @@ async def main_streaming(args):
t_prev = t_start
chunk_idx = 0
async for stage_output in omni.generate(prompt, request_id=request_id):
mm = stage_output.request_output.outputs[0].multimodal_output
mm = stage_output.outputs[0].multimodal_output
if not stage_output.finished:
t_now = time.perf_counter()
audio = mm.get("audio")
@@ -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",
@@ -116,7 +116,7 @@ def main(args: Any) -> None:
prompt = _build_custom_voice_input(args)
final_output = None
for stage_outputs in omni.generate([prompt]):
final_output = stage_outputs.request_output
final_output = stage_outputs
if final_output is None:
raise RuntimeError("Qwen3-TTS did not produce an output.")
@@ -38,7 +38,7 @@ def _stage_durations_from_output(outputs) -> dict[str, float]:
durations = getattr(omni_out, "stage_durations", None) or {}
if durations:
return dict(durations)
ro = getattr(omni_out, "request_output", None)
ro = omni_out
if ro is not None:
return dict(getattr(ro, "stage_durations", None) or {})
return {}
@@ -99,7 +99,7 @@ def extract_audio(outputs) -> tuple[np.ndarray, int]:
out = outputs[0]
mm = getattr(out, "multimodal_output", None)
if mm is None:
ro = getattr(out, "request_output", None)
ro = out
if ro is not None:
mm = getattr(ro, "multimodal_output", None)
if mm is None and ro.outputs:
@@ -8,6 +8,7 @@ A unified script for text-to-video generation. Supports multiple models with mod
|---|---|---|---|---|---|
| `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) |
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) |
| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
@@ -57,18 +58,31 @@ python text_to_video.py \
### LingBot-Video
The dedicated LingBot example explicitly selects the video output modality:
The shared runner recognizes the official dense and MoE checkpoint IDs, selects
`LingBotVideoPipeline`, and builds the canonical video request envelope:
```bash
python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
python text_to_video.py \
--model robbyant/lingbot-video-dense-1.3b \
--prompt "a robotic arm picks up a red block" \
--height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
--guidance-scale 3.0 --flow-shift 3.0 --fps 24 \
--output lingbot_t2v.mp4
```
Use the same command with `robbyant/lingbot-video-moe-30b-a3b` for the MoE
checkpoint. Requested frame counts are rounded upward to the causal
VAE's `4n+1` grid.
VAE's `4n+1` grid. For a local checkpoint path whose name does not contain
`lingbot`, pass `--model-class-name LingBotVideoPipeline`.
LingBot-only options such as `batch_cfg` and `output_type` use the shared
model-extra channel, for example:
```bash
python text_to_video.py \
--model robbyant/lingbot-video-dense-1.3b \
--extra-body '{"batch_cfg": true, "output_type": "np"}'
```
### LTX-2
@@ -16,7 +16,7 @@ from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.lora.request import LoRARequest
from vllm_omni.lora.utils import stable_lora_int_id
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name, get_output_tensor_range
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
@@ -80,6 +80,17 @@ _MODEL_PRESETS = {
"fps": 16,
"output": "helios_output.mp4",
},
"lingbot": {
"model_class_name": "LingBotVideoPipeline",
"height": 192,
"width": 320,
"num_frames": 9,
"num_inference_steps": 2,
"guidance_scale": 3.0,
"fps": 24,
"flow_shift": 3.0,
"output": "lingbot_video_output.mp4",
},
"ltx2": {
"height": 512,
"width": 768,
@@ -110,27 +121,51 @@ _MODEL_PRESETS = {
def _detect_preset(model: str, model_class_name: str | None = None) -> dict:
model_lower = model.lower()
class_lower = (model_class_name or "").lower()
if "distilled" in class_lower or "distilled" in model_lower:
return _MODEL_PRESETS["ltx2_distilled"]
if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
return _MODEL_PRESETS["ltx23"]
if "ltx2" in class_lower or "ltx-2" in model_lower or "ltx_2" in model_lower:
if "lingbot" in model_lower or "lingbotvideo" in class_lower:
return _MODEL_PRESETS["lingbot"]
if "ltx" in class_lower or "ltx" in model_lower:
if "distilled" in class_lower or "distilled" in model_lower:
return _MODEL_PRESETS["ltx2_distilled"]
if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
return _MODEL_PRESETS["ltx23"]
return _MODEL_PRESETS["ltx2"]
if "vace" in model_lower:
if "vace" in model_lower or "vace" in class_lower:
return _MODEL_PRESETS["vace"]
# Edge must be matched before the generic cosmos branch (its "cosmos" substring would
# otherwise pick up the Nano/Super 720p / gs 6.0 / flow_shift 10.0 preset).
if ("cosmos" in model_lower or "cosmos" in class_lower) and ("edge" in model_lower or "edge" in class_lower):
return _MODEL_PRESETS["cosmos3_edge"]
if "cosmos" in model_lower:
if "cosmos" in model_lower or "cosmos" in class_lower:
return _MODEL_PRESETS["cosmos"]
if "hunyuan" in model_lower:
if "hunyuan" in model_lower or "hunyuan" in class_lower:
return _MODEL_PRESETS["hunyuan"]
if "helios" in model_lower:
if "helios" in model_lower or "helios" in class_lower:
return _MODEL_PRESETS["helios"]
return _MODEL_PRESETS["wan"]
def build_text_to_video_prompt(prompt: str, negative_prompt: str | None) -> dict[str, Any]:
"""Build the canonical request envelope for the shared T2V example."""
result: dict[str, Any] = {
"prompt": prompt,
"modalities": ["video"],
}
if negative_prompt is not None:
result["negative_prompt"] = negative_prompt
return result
def _normalize_float_tensor(tensor: torch.Tensor, source_range: str) -> torch.Tensor:
"""Normalize decoded frames according to the pipeline's output contract."""
if not tensor.is_floating_point():
return tensor
if source_range == "negative_one_to_one":
return tensor.clamp(-1, 1) * 0.5 + 0.5
if source_range == "zero_to_one":
return tensor.clamp(0, 1)
raise ValueError(f"Unsupported floating-point tensor range: {source_range!r}")
def parse_profiler_config(value: str) -> dict[str, Any]:
try:
config = json.loads(value)
@@ -160,7 +195,7 @@ def parse_extra_body(value: str) -> dict[str, Any]:
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Generate a video from a text prompt. "
"Supports Wan2.2, HunyuanVideo-1.5, Helios, and other text-to-video models."
"Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, and other text-to-video models."
)
parser.add_argument(
"--model",
@@ -390,7 +425,7 @@ def _extract_peak_memory_mb(result: Any) -> float:
return 0.0
val = getattr(result, "peak_memory_mb", 0.0)
if not val:
inner = getattr(result, "request_output", None)
inner = result
if isinstance(inner, list):
inner = inner[0] if inner else None
val = getattr(inner, "peak_memory_mb", 0.0)
@@ -408,6 +443,7 @@ def main():
for key, default_val in preset.items():
if getattr(args, key.replace("-", "_"), None) is None:
setattr(args, key.replace("-", "_"), default_val)
model_class_name = args.model_class_name
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
# Cache-dit config (Wan2.2 only)
@@ -480,6 +516,7 @@ def main():
omni = Omni(**omni_kwargs)
model_class_name = get_model_class_name(omni) or model_class_name
declared_extra_body_params = get_extra_body_params(model_class_name)
output_tensor_range = get_output_tensor_range(model_class_name)
if profiler_enabled:
print("[Profiler] Starting profiling...")
@@ -513,12 +550,11 @@ def main():
lora_path=lora_path,
)
prompt_dict = {"prompt": args.prompt}
if args.negative_prompt is not None:
prompt_dict["negative_prompt"] = args.negative_prompt
elif preset not in (_MODEL_PRESETS["ltx2"], _MODEL_PRESETS["ltx23"]):
negative_prompt = args.negative_prompt
if negative_prompt is None and all(preset is not _MODEL_PRESETS[name] for name in ("lingbot", "ltx2", "ltx23")):
# Preserve the historical empty-prompt behavior for non-LTX examples.
prompt_dict["negative_prompt"] = ""
negative_prompt = ""
prompt_dict = build_text_to_video_prompt(args.prompt, negative_prompt)
extra_args = {}
if lora_request:
@@ -575,8 +611,8 @@ def main():
if frames.multimodal_output and "audio" in frames.multimodal_output:
audio = frames.multimodal_output["audio"]
audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
if frames.is_pipeline_output and frames.request_output is not None:
inner_output = frames.request_output
if frames.is_pipeline_output and frames is not None:
inner_output = frames
if isinstance(inner_output, OmniRequestOutput):
if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
audio = inner_output.multimodal_output["audio"]
@@ -630,8 +666,7 @@ def main():
frame_tensor = frame_tensor[0]
if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4):
frame_tensor = frame_tensor.permute(1, 2, 0)
if frame_tensor.is_floating_point():
frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5
frame_tensor = _normalize_float_tensor(frame_tensor, output_tensor_range)
return frame_tensor.float().numpy()
if isinstance(frame, np.ndarray):
frame_array = frame
@@ -683,8 +718,7 @@ def main():
video_tensor = video_tensor[0]
elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
video_tensor = video_tensor.permute(1, 2, 3, 0)
if video_tensor.is_floating_point():
video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
video_tensor = _normalize_float_tensor(video_tensor, output_tensor_range)
video_array = video_tensor.float().numpy()
elif isinstance(frames, np.ndarray):
video_array = frames
@@ -1,131 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import argparse
import time
from pathlib import Path
from typing import Any
import numpy as np
import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate text-to-video with a LingBot-Video checkpoint.")
parser.add_argument("--model", default="robbyant/lingbot-video-dense-1.3b")
parser.add_argument("--prompt", default="a robotic arm picks up a red block")
parser.add_argument("--negative-prompt", default=None)
parser.add_argument("--output", default="lingbot_video_output.mp4")
parser.add_argument("--height", type=int, default=192)
parser.add_argument("--width", type=int, default=320)
parser.add_argument("--num-frames", type=int, default=9)
parser.add_argument("--num-inference-steps", type=int, default=2)
parser.add_argument("--guidance-scale", type=float, default=3.0)
parser.add_argument("--flow-shift", type=float, default=3.0)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--fps", type=int, default=24)
parser.add_argument("--output-type", choices=["pt", "np"], default="pt")
parser.add_argument("--batch-cfg", action="store_true")
return parser.parse_args()
def _peak_memory_mb(output: Any) -> float:
result = output[0] if isinstance(output, list) and output else output
return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
def _extract_video(output: Any) -> Any:
if isinstance(output, list):
output = output[0] if output else None
if isinstance(output, OmniRequestOutput):
if output.final_output_type != "image":
raise ValueError(f"Unexpected output type {output.final_output_type!r}.")
if not output.images:
raise ValueError("No video frames found in OmniRequestOutput.")
return output.images[0] if len(output.images) == 1 else output.images
return output
def _to_frame_array(video: Any) -> np.ndarray:
if isinstance(video, torch.Tensor):
tensor = video.detach().cpu()
if tensor.dim() == 5:
tensor = tensor[0]
if tensor.dim() == 4 and tensor.shape[-1] not in (3, 4):
if tensor.shape[0] in (3, 4):
tensor = tensor.permute(1, 2, 3, 0)
elif tensor.shape[1] in (3, 4):
tensor = tensor.permute(0, 2, 3, 1)
if tensor.is_floating_point():
if float(tensor.min()) < 0.0:
tensor = tensor.clamp(-1, 1) * 0.5 + 0.5
else:
tensor = tensor.clamp(0, 1)
return tensor.float().numpy()
if isinstance(video, np.ndarray):
array = video[0] if video.ndim == 5 else video
if np.issubdtype(array.dtype, np.integer):
return array.astype(np.float32) / 255.0
return array.astype(np.float32, copy=False)
if isinstance(video, list):
return np.stack([_to_frame_array(frame) for frame in video], axis=0)
return np.asarray(video, dtype=np.float32)
def main() -> None:
args = parse_args()
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
omni = Omni(
model=args.model,
model_class_name="LingBotVideoPipeline",
flow_shift=args.flow_shift,
parallel_config=DiffusionParallelConfig(),
)
prompt: dict[str, Any] = {
"prompt": args.prompt,
"modalities": ["video"],
"num_frames": args.num_frames,
}
if args.negative_prompt is not None:
prompt["negative_prompt"] = args.negative_prompt
sampling_params = OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
num_frames=args.num_frames,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
generator=generator,
output_type=args.output_type,
extra_args={"batch_cfg": args.batch_cfg, "shift": args.flow_shift},
)
start = time.perf_counter()
output = omni.generate(prompt, sampling_params)
elapsed = time.perf_counter() - start
print(f"Total generation time: {elapsed:.4f} seconds ({elapsed * 1000:.2f} ms)")
peak_mb = _peak_memory_mb(output)
if peak_mb:
print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
video = _to_frame_array(_extract_video(output))
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
from diffusers.utils import export_to_video
export_to_video(list(video), str(output_path), fps=args.fps)
print(f"Saved generated video to {output_path}")
omni.close()
if __name__ == "__main__":
main()
@@ -138,7 +138,7 @@ def _extract_peak_memory_mb(result: Any) -> float:
return 0.0
val = getattr(result, "peak_memory_mb", 0.0)
if not val:
inner = getattr(result, "request_output", None)
inner = result
if isinstance(inner, list):
inner = inner[0] if inner else None
val = getattr(inner, "peak_memory_mb", 0.0)
@@ -605,6 +605,13 @@ python qwen3_tts/streaming_speech_client.py \
The client writes one PCM file per sentence and a matching
`sentence_XXX_timestamps.json` sidecar.
Non-streaming requests can also ask for timestamps: pass
`"word_timestamps": true` to `POST /v1/audio/speech` and read the
`X-Word-Timestamps` response header (JSON list of `{word, start_ms, end_ms}`,
ASCII-escaped). Past 4 KB the header is replaced by
`X-Word-Timestamps-Omitted: oversize; bytes=<n>; limit=4096` and the audio
still returns — use the WebSocket path for long transcripts.
To *see* the alignment instead of reading a JSON sidecar, run the
word-timestamp Gradio demo (server must be launched with `--forced-aligner`):
```bash
+3
View File
@@ -73,6 +73,8 @@ 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
@@ -111,6 +113,7 @@ plugins:
- "vllm_omni.entrypoints.openai" # avoid importing vllm in mkdocs building
- "vllm_omni.entrypoints.openai.protocol" # avoid importing vllm in mkdocs building
- "vllm_omni.experimental" # avoid importing optional serving deps (openai, fastapi, uvicorn, httpx) in mkdocs building
- "vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored" # vendored NeMo code; upstream docstrings trip griffe strict warnings
- mkdocstrings:
handlers:
python:
+13
View File
@@ -206,6 +206,12 @@ ignore = [
# These files keep upstream conventions (star imports, long lines, math-symbol
# argument names) so they stay easy to diff against upstream releases.
"vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py" = ["E402", "E501", "F403", "F405", "N803"]
# NemotronVoiceChat inference modules vendored from NVIDIA-NeMo/Speech (branch
# nemotron-labs-voicechat). Numerics-bearing code is kept byte-close to
# upstream to stay diffable and parity-safe, so upstream style (long lines,
# legacy typing, whitespace, unused-import shims) is exempted. TID251 and the
# F-correctness rules still apply.
"vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/**" = ["E501", "E721", "E741", "F401", "F841", "I001", "N803", "N814", "UP006", "UP007", "UP008", "UP011", "UP032", "UP035", "W291", "W293"]
"vllm_omni/experimental/fullduplex/joyvl/decision/prompts.py" = ["E501"] # System/summarizer prompts are long prose lines
[tool.ruff.lint.flake8-tidy-imports.banned-api]
@@ -314,6 +320,10 @@ extend-ignore-identifiers-re = [
".*[Nn][Dd]3.*",
".*updator.*",
"byt5.*",
# NemotronVoiceChat vendored triton kernel: dout/DOUT abbreviate d_out
# (output feature dim), not "doubt".
".*dout.*",
".*DOUT.*",
]
extend-ignore-re = [
"\\.\\.\\.n,d->\\.\\.\\.nd",
@@ -334,6 +344,9 @@ byt5 = "byt5"
# "seeked" is the standard HTML media element DOM event name (fires when a
# seek completes), used by the word-timestamps Gradio demo. Not a typo.
seeked = "seeked"
# Ryan Leary: author credited in the NemotronVoiceChat vendored NeMo
# features.py copyright header. A surname, not a typo of "leery".
Leary = "Leary"
# IndexTTS 2.5 language codes and speech-emotion-recognition event token.
fo = "fo"
ba = "ba"
+113
View File
@@ -0,0 +1,113 @@
# NemotronLabs VoiceChat 11B
> Offline speech-to-speech: a 16 kHz user utterance in, the agent's spoken reply
> (text + 22.05 kHz WAV) out, on a 3-stage vLLM-Omni pipeline.
## Summary
- Vendor: NVIDIA
- Model: [`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
- Task: Speech-to-speech voice chat. The model runs a frame-locked 12.5 Hz
timeline: a Conformer + NemotronH hybrid-Mamba thinker emits one text token
per acoustic frame, a Gemma3-1B EAR-TTS talker turns the text timeline into
31-quantizer RVQ code stacks, and an RVQ-VAE codec decodes them to audio.
- Mode: Offline single-turn inference (batch=1). Online/duplex serving,
function calling, and batch>1 are documented follow-ups.
- Maintainer: [`@yuekaizhang`](https://github.com/yuekaizhang)
## When to use this recipe
Use this recipe to run a single-turn voice-chat exchange with
`NVIDIA-NemotronLabs-VoiceChat-11B` on one GPU: you provide a user utterance as
a WAV file (any sample rate; it is resampled to 16 kHz mono) plus an optional
spoken-style system prompt, and get back the agent's reply as text and a
22.05 kHz WAV. The integration is NeMo-free at runtime — the perception
Conformer, EAR-TTS talker, and RVQ-VAE codec are vendored, dependency-stripped
NeMo modules (`nemo_vendored/`), so no `nemo_toolkit` install is needed.
## References
- Offline example:
[`examples/offline_inference/nemotron_voicechat/end2end.py`](../../examples/offline_inference/nemotron_voicechat/end2end.py)
- Model modules (thinker / talker / code2wav / vendored NeMo):
[`vllm_omni/model_executor/models/nemotron_voicechat/`](../../vllm_omni/model_executor/models/nemotron_voicechat/)
- Staged pipeline config:
[`vllm_omni/deploy/nemotron_labs_voicechat.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat.yaml)
- Upstream model card:
[`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
- Reference implementation: NVIDIA-NeMo/Speech, branch `nemotron-labs-voicechat`
## Pipeline
| stage | arch | dtype | role |
|---|---|---|---|
| 0 thinker | `NemotronVoiceChatThinkerForConditionalGeneration` (LLM_AR) | fp32 | WAV + system prompt -> frame-locked text-token timeline (+ function channel) |
| 1 talker | `NemotronVoiceChatTalker` (LLM_AR) | fp32 | text timeline -> 31-quantizer RVQ code stacks (one per 80 ms frame) |
| 2 code2wav | `NemotronVoiceChatCode2Wav` (LLM_GENERATION) | fp32 | RVQ-VAE decode -> 22.05 kHz PCM |
Stages 0/1 default to fp32 for exact parity with the NeMo reference
implementation (greedy decoding matches it token for token on the acceptance
fixture). The deploy yaml documents a ~2x-faster bf16 thinker option whose
output stayed within one word of the reference in testing.
## Hardware Support
## GPU
### 1x H100 80GB
#### Environment
- OS: Linux
- Python: 3.12
- vLLM version: 0.26.0
- vLLM-Omni version or commit: this PR / current `main`
#### Command
```bash
# Tokenizer: the checkpoint ships no HF tokenizer; it resolves from the
# nvidia/NVIDIA-Nemotron-Nano-9B-v2 HF id automatically. For air-gapped runs,
# point NEMOTRON_VOICECHAT_LLM_PATH at a local snapshot of that repo instead.
python examples/offline_inference/nemotron_voicechat/end2end.py \
--checkpoint /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
--wav /path/to/user_question.wav \
--output-dir results/nemotron_voicechat
```
#### Verification
```bash
ls results/nemotron_voicechat
# <stem>_output.txt the agent reply as text
# <stem>_output.wav the agent reply as 22.05 kHz audio
# <stem>_text_tokens.json the frame-locked text-token timeline
```
The reply text should read as a coherent spoken-style answer to the question in
the input WAV, and the WAV should transcribe to (approximately) the same text
with any ASR model.
#### Notes
- Memory usage: the shipped yaml runs all three stages on one GPU
(`gpu_memory_utilization` 0.62 / 0.12 / 0.06); peak usage is dominated by the
fp32 thinker. The fp32 default has a hard floor of roughly 43 GB of thinker
weights alone (9B backbone + 587M `embed_tokens` + 587M `function_head` +
0.6B Conformer), so 48 GB cards cannot run it — use the bf16 thinker option
documented in the deploy yaml on anything smaller than an 80 GB part.
- Input sizing: the timeline is frame-locked, so the reply budget IS the input
duration. The acoustic channel trails the text channel; if the WAV does not
carry enough trailing silence for the reply to finish, the spoken answer is
truncated silently. Leave generous trailing silence (a question ending at
~4.5 s truncated in an 8 s WAV but completed cleanly in 16 s); the offline
example warns when the text channel is still speaking near the last frame.
- Key flags: sampling is greedy end to end. The thinker is frame-locked —
`max_tokens` equals the acoustic frame count with `ignore_eos=True`. Do NOT
set `min_tokens` on the thinker: the tokenizer's EOS token is also the
frame-locked PAD/silence token, so masking it forces the model to babble
instead of pausing.
- The talker's `max_tokens` is 16383 (its stage prompt is one placeholder
token, and the stage context is 16384).
- Known limitations: batch=1 only; offline single turn (no online/duplex
serving yet); the function-call channel is decoded but not yet acted on.
+1
View File
@@ -66,6 +66,7 @@ recipes/
| [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
| [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
| [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB |
| [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB |
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
| [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
+3 -3
View File
@@ -27,7 +27,7 @@ parallel, cache, quantized, or expert-kernel backends.
- Dense checkpoint: <https://huggingface.co/robbyant/lingbot-video-dense-1.3b>
- MoE checkpoint: <https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b>
- Upstream project: <https://github.com/Robbyant/lingbot-video>
- Related offline example: [`examples/offline_inference/text_to_video/text_to_video_lingbot.py`](../../examples/offline_inference/text_to_video/text_to_video_lingbot.py)
- Related offline example: [`examples/offline_inference/text_to_video/text_to_video.py`](../../examples/offline_inference/text_to_video/text_to_video.py)
- Related online video API docs: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)
## Hardware Support
@@ -52,7 +52,7 @@ The MoE path is validated with BF16 expert weights.
```bash
CUDA_VISIBLE_DEVICES=0 \
python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
python examples/offline_inference/text_to_video/text_to_video.py \
--model robbyant/lingbot-video-dense-1.3b \
--prompt "a robotic arm picks up a red block" \
--output lingbot_t2v.mp4 \
@@ -70,7 +70,7 @@ python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
```bash
CUDA_VISIBLE_DEVICES=0 \
python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
python examples/offline_inference/text_to_video/text_to_video.py \
--model robbyant/lingbot-video-moe-30b-a3b \
--prompt "a robotic arm picks up a red block" \
--output lingbot_moe_t2v.mp4 \
+1 -1
View File
@@ -82,7 +82,7 @@ example and the DiT sections below.
- Optional environment variables:
```bash
export CACHE_DIT_VERSION=1.3.0
export CACHE_DIT_VERSION=1.5.0
```
HunyuanImage-3.0 sets the diffusion attention backend to `TORCH_SDPA`
+3 -1
View File
@@ -16,7 +16,7 @@ diffusers==0.38.0
safetensors>=0.8.0
accelerate==1.12.0
soundfile>=0.13.1
cache-dit==1.3.0
cache-dit==1.5.0
tqdm>=4.66.0
torchsde>=0.2.6
openai-whisper>=20250625
@@ -34,3 +34,5 @@ pydantic-settings>=2.0.0
# base image does not ship it, so declare it explicitly. The loader also guards
# this import so a missing gguf no longer breaks all diffusion model loading.
gguf>=0.10.0
# Required for running the Cosmos3 model with guardrails enabled (default behavior).
cosmos-guardrail>=0.3.1
@@ -0,0 +1,92 @@
# SPDX-License-Identifier: Apache-2.0
"""Offline Daily-Omni sources: local mirrors must not fall back to the Hub id."""
import json
import os
import tarfile
from pathlib import Path
import pytest
from vllm_omni.benchmarks.data_modules.daily_omni_dataset import (
daily_omni_local_qa_json,
daily_omni_local_videos_dir,
resolve_daily_omni_local_root,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
def _write_qa_json(root: Path) -> Path:
qa = root / "qa.json"
qa.write_text(json.dumps([{"Question": "q", "Answer": "A", "video_id": "vid0"}]), encoding="utf-8")
return qa
def test_resolve_local_root_returns_none_for_hub_id() -> None:
assert resolve_daily_omni_local_root("liarliar/Daily-Omni") is None
assert resolve_daily_omni_local_root("") is None
assert resolve_daily_omni_local_root(None) is None
def test_plain_local_mirror(tmp_path: Path) -> None:
_write_qa_json(tmp_path)
(tmp_path / "Videos" / "vid0").mkdir(parents=True)
root = resolve_daily_omni_local_root(str(tmp_path))
assert root == tmp_path.resolve()
assert daily_omni_local_qa_json(root) == (tmp_path / "qa.json").resolve()
assert daily_omni_local_videos_dir(root) == root / "Videos"
def test_hf_cache_dir_resolves_to_snapshot(tmp_path: Path) -> None:
cache_dir = tmp_path / "datasets--liarliar--Daily-Omni"
snapshot = cache_dir / "snapshots" / "deadbeef"
snapshot.mkdir(parents=True)
_write_qa_json(snapshot)
(cache_dir / "refs").mkdir()
(cache_dir / "refs" / "main").write_text("deadbeef", encoding="utf-8")
root = resolve_daily_omni_local_root(str(cache_dir))
assert root == snapshot.resolve()
assert daily_omni_local_qa_json(root) == (snapshot / "qa.json").resolve()
def test_hf_cache_dir_without_refs_picks_newest_snapshot(tmp_path: Path) -> None:
cache_dir = tmp_path / "datasets--liarliar--Daily-Omni"
older = cache_dir / "snapshots" / "aaa"
newer = cache_dir / "snapshots" / "bbb"
older.mkdir(parents=True)
newer.mkdir(parents=True)
os.utime(older, (1, 1))
os.utime(newer, (10, 10))
assert resolve_daily_omni_local_root(str(cache_dir)) == newer.resolve()
def test_videos_tar_is_extracted(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("HF_HOME", str(tmp_path / "hf_home"))
root = tmp_path / "mirror"
(root / "Videos" / "vid0").mkdir(parents=True)
(root / "Videos" / "vid0" / "vid0_video.mp4").write_bytes(b"fake")
with tarfile.open(root / "Videos.tar", "w") as tf:
tf.add(root / "Videos", arcname="Videos")
# Only the tarball is left in the mirror.
for child in (root / "Videos").rglob("*"):
if child.is_file():
child.unlink()
(root / "Videos" / "vid0").rmdir()
(root / "Videos").rmdir()
videos_dir = daily_omni_local_videos_dir(root)
assert videos_dir is not None
assert (videos_dir / "vid0" / "vid0_video.mp4").is_file()
# Second call reuses the extraction instead of redoing it.
assert daily_omni_local_videos_dir(root) == videos_dir
def test_missing_qa_json_reports_none(tmp_path: Path) -> None:
root = resolve_daily_omni_local_root(str(tmp_path))
assert daily_omni_local_qa_json(root) is None
assert daily_omni_local_videos_dir(root) is None
+7 -19
View File
@@ -2,9 +2,8 @@
The ``minicpm-interleave`` mode is an accuracy-critical protocol: MiniCPM-o only reaches
OpenBMB's reported Daily-Omni score when frames and 1s audio segments arrive strictly
interleaved, with ``max_slice_nums=1`` / ``use_image_id=False`` and an empty MiniCPM
system message (not the Qwen system prompt). These tests pin that contract so a refactor
cannot silently regress the score.
interleaved, with ``max_slice_nums=1`` / ``use_image_id=False`` and no Qwen system prompt.
These tests pin that contract so a refactor cannot silently regress the score.
vllm stubs are installed by tests/benchmarks/conftest.py before collection.
"""
@@ -99,22 +98,13 @@ def _stub_extract(monkeypatch, dataset: DailyOmniDataset, counter: list[int]) ->
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(("n", "k"), [(10, 4), (100, 64), (1000, 64), (7, 7), (3, 5), (1, 1)])
def test_uniform_sample_indices_matches_openbmb_midpoint(n: int, k: int) -> None:
"""Pin OpenBMB midpoint-bin sampling (not np.linspace endpoints)."""
@pytest.mark.parametrize(("n", "k"), [(10, 4), (100, 64), (7, 7), (3, 5), (1, 1)])
def test_uniform_sample_indices_matches_numpy_linspace(n: int, k: int) -> None:
got = _uniform_sample_indices(n, k)
if n <= k:
assert got == list(range(n))
else:
gap = n / k
expected = [int(i * gap + gap / 2) for i in range(k)]
assert got == expected
# Endpoints differ from linspace for typical Daily-Omni downsample sizes.
linspace = [int(i) for i in np.linspace(0, n - 1, k, dtype=int)]
if n == 1000 and k == 64:
assert got[0] == 7 and got[-1] == 992
assert linspace[0] == 0 and linspace[-1] == 999
assert got != linspace
assert got == [int(i) for i in np.linspace(0, n - 1, k, dtype=int)]
def test_uniform_sample_indices_non_positive_k() -> None:
@@ -223,16 +213,14 @@ def test_minicpm_interleave_missing_local_video_returns_none(qa_json, tmp_path)
# ---------------------------------------------------------------------------
def test_minicpm_interleave_keeps_empty_system_message(monkeypatch, qa_json, video_dir) -> None:
"""Official MiniCPM omni mode still sends role=system with empty content."""
def test_minicpm_interleave_omits_qwen_system_message(monkeypatch, qa_json, video_dir) -> None:
ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all")
_stub_extract(monkeypatch, ds, [])
parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None)
messages = ds._build_daily_omni_openai_messages(parts, "q?", {"A": "a"})
assert [m["role"] for m in messages] == ["system", "user"]
assert messages[0]["content"][0]["text"] == ""
assert [m["role"] for m in messages] == ["user"]
def test_qwen_pack_mode_keeps_system_message(qa_json, video_dir) -> None:
+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``)."""
@@ -0,0 +1,71 @@
import pytest
import vllm_omni.config.pipeline_registry # noqa: F401 (populate registry)
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES as _PIPELINE_REGISTRY
from vllm_omni.config.stage_config import (
DeployConfig,
StageExecutionType,
)
from vllm_omni.model_executor.stage_input_processors.forced_aligner import (
POOLING_OUTPUT_DECODER_PATH,
TIMESTAMPS_MODALITY,
)
from vllm_omni.utils.forced_aligner import _DEFAULT_CONFIG_PATH, inject_forced_aligner_stage
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_inject_appends_pooling_aligner_stage():
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
deploy = DeployConfig()
n0 = len(pipeline.stages)
ext_pipeline, ext_deploy = inject_forced_aligner_stage(
pipeline, deploy, {"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"}
)
# A new pooling stage is appended; the original pipeline is unchanged.
assert len(pipeline.stages) == n0
assert len(ext_pipeline.stages) == n0 + 1
# The original deploy config is not mutated either (double-injection safe).
assert len(deploy.stages) == 0
assert len(ext_deploy.stages) == 1
aligner = ext_pipeline.stages[-1]
# Pooling stage reuses LLM_AR (+ runner=pooling), no dedicated LLM_POOLING.
assert aligner.execution_type == StageExecutionType.LLM_AR
assert aligner.final_output is True
assert aligner.final_output_type == TIMESTAMPS_MODALITY
assert aligner.input_sources == (n0 - 1,)
assert aligner.requires_multimodal_data is True
assert aligner.owns_tokenizer is True
# The deploy stage carries the separate aligner checkpoint + decoder hook.
ds = ext_deploy.stages[-1]
assert ds.engine_extras["model"] == "/models/Qwen3-ForcedAligner-0.6B"
assert ds.engine_extras["runner"] == "pooling"
assert ds.engine_extras["pooling_output_decoder"] == POOLING_OUTPUT_DECODER_PATH
def test_inject_from_config_only():
# --forced-aligner-config alone (model in the YAML) must also inject.
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
deploy = DeployConfig()
n0 = len(pipeline.stages)
ext_pipeline, ext_deploy = inject_forced_aligner_stage(
pipeline, deploy, {"forced_aligner": None, "forced_aligner_config": str(_DEFAULT_CONFIG_PATH)}
)
assert len(ext_pipeline.stages) == n0 + 1
assert ext_pipeline.stages[-1].final_output_type == TIMESTAMPS_MODALITY
assert ext_deploy.stages[-1].engine_extras["model"] == "Qwen/Qwen3-ForcedAligner-0.6B"
def test_inject_noop_without_forced_aligner():
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
deploy = DeployConfig()
ext_pipeline, ext_deploy = inject_forced_aligner_stage(pipeline, deploy, {})
assert ext_pipeline is pipeline
assert len(ext_deploy.stages) == 0
+5
View File
@@ -513,6 +513,7 @@ def test_sub_config_fields_match_structured_scopes():
"enable_prefix_caching",
"disable_hybrid_kv_cache_manager",
"mm_processor_cache_gb",
"mamba_ssm_cache_dtype",
}
assert {f.name for f in fields(OmniStageSchedulerConfig)} == {
"max_num_seqs",
@@ -581,8 +582,12 @@ def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
assert cfg.data_parallel_size == 3
assert cfg.cfg_parallel_size == 3
assert cfg.mask_sp_padding is True
assert cfg.vae_parallel_mode == "tile"
assert cfg.world_size == 72
auto_cfg = OmniStageDiffusionParallelConfig(vae_parallel_mode="auto")
assert auto_cfg.vae_parallel_mode == "auto"
def test_parallel_config_derived_fields_are_not_init_inputs():
with pytest.raises(ValidationError):
@@ -27,6 +27,7 @@ _MIXIN_UPDATE_HELPERS = (
"_attach_finished_request_sets",
"_attach_scheduler_stats",
"_capture_omni_connector_output",
"_maybe_decode_pooling_output",
)
@@ -138,6 +139,7 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
pending_stop_after_extraction=set(),
waiting_for_transfer_free=set(),
_new_prompt_len_snapshot={},
_pooling_output_decoder=None,
finished_req_ids=set(),
finished_req_ids_dict=defaultdict(set),
kv_cache_manager=SimpleNamespace(take_events=lambda: None),
@@ -267,3 +269,72 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
assert output_by_id["bad"].new_token_ids == []
assert output_by_id["good"].new_token_ids == [8]
np.testing.assert_array_equal(output_by_id["good"].new_logprobs.logprob_token_ids[:, 0], [8])
def _pooling_model_runner_output(pooler_tensor) -> SimpleNamespace:
return SimpleNamespace(
sampled_token_ids=[[]],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[pooler_tensor],
num_nans_in_logits=None,
kv_connector_output=None,
cudagraph_stats=None,
req_id_to_index={"req": 0},
routed_experts=None,
)
def _make_pooling_stub(decoder) -> tuple[_Request, SimpleNamespace]:
request = _Request("req")
request.sampling_params = SimpleNamespace(num_logprobs=None)
request.pooling_params = SimpleNamespace()
scheduler = _make_scheduler_stub([request])
scheduler._pooling_output_decoder = decoder
scheduler.vllm_config = SimpleNamespace(model_config=SimpleNamespace(hf_config=None))
_bind_request_lifecycle(scheduler, update_request=lambda req, tokens: (tokens, False))
return request, scheduler
def test_pooling_decode_success_emits_decoded_payload() -> None:
import torch
decoded = torch.tensor([[1, 2]], dtype=torch.int32)
request, scheduler = _make_pooling_stub(lambda logits, req, hf: decoded)
scheduler_output = SimpleNamespace(
num_scheduled_tokens={"req": 1},
scheduled_spec_decode_tokens={},
num_invalid_spec_tokens=0,
)
outputs = OmniARScheduler.update_from_output(
scheduler, scheduler_output, _pooling_model_runner_output(torch.zeros(2, 4))
)
(output,) = outputs[0].outputs
assert output.finish_reason is FinishReason.STOP
assert output.pooling_output is decoded
def test_pooling_decode_failure_finishes_request_with_error() -> None:
import torch
def raising_decoder(logits, req, hf):
raise ValueError("boom")
request, scheduler = _make_pooling_stub(raising_decoder)
scheduler_output = SimpleNamespace(
num_scheduled_tokens={"req": 1},
scheduled_spec_decode_tokens={},
num_invalid_spec_tokens=0,
)
outputs = OmniARScheduler.update_from_output(
scheduler, scheduler_output, _pooling_model_runner_output(torch.zeros(2, 4))
)
(output,) = outputs[0].outputs
# Decoder failure is a request error (500), not an empty success.
assert output.finish_reason is FinishReason.ERROR
assert output.pooling_output is None
assert "pooling output decode failed" in str(request.stop_reason)
@@ -120,6 +120,9 @@ class TestFullPayloadCoordinatorSelection(unittest.TestCase):
("MiMoAudioModel", "code2wav"),
("Qwen3TTSCode2Wav", "code2wav"),
("MiniCPMO45Code2Wav", "code2wav"),
# nemotron_voicechat: sync-path RVQ-VAE decode waits for the
# talker's full [frames, 31] code-stack payload.
("NemotronVoiceChatCode2Wav", "code2wav"),
("CosyVoice3Model", "cosyvoice3_code2wav"),
# audex: sync-path TTS decoder (async_chunk: false deploys) and
# the TTA XCodec1 stage (always sync full-payload).
@@ -0,0 +1,80 @@
[
{
"test_name": "test_minicpmo_4_5_ready",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100",
"npu": "A3"
},
"num_cards": 1
}
},
"core_model",
"omni"
],
"server_params": {
"model": "openbmb/MiniCPM-o-4_5",
"extra_cli_args": [
"--deploy-config",
"vllm_omni/deploy/minicpmo_4_5.yaml",
"--trust-remote-code"
]
},
"benchmark_params": [
{
"dataset_name": "seed-tts",
"dataset_path": "zhaochenyang20/seed-tts-eval",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [
10,
40
],
"max_concurrency": [
1,
4
],
"no_oversample": true,
"trust_remote_code": true,
"seed_tts_locale": "en",
"extra_body": {
"modalities": [
"text",
"audio"
],
"chat_template_kwargs": {
"enable_thinking": false,
"use_tts_template": true
}
},
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration",
"baseline": {
"H100": {
"request_throughput": [
0.7803,
1.1008
],
"mean_ttft_ms": [
166.9764,
307.7133
],
"mean_e2el_ms": [
1361.1682,
3800.6332
],
"mean_audio_ttfp_ms": [
635.8099,
1615.6591
],
"mean_audio_rtf": [
0.3151,
0.8568
]
}
}
}
]
}
]
@@ -125,7 +125,7 @@ def _extract_images(output: OmniRequestOutput) -> list:
return output.images
# When the output comes from the orchestrator pipeline, images may be
# nested inside request_output.
inner = getattr(output, "request_output", None)
inner = output
if inner is not None and hasattr(inner, "images") and inner.images:
return inner.images
return []
+2 -2
View File
@@ -56,10 +56,10 @@ def test_teacache(model_name: str):
# Extract images from request_output['images']
first_output = outputs[0]
assert first_output.final_output_type == "image"
if not hasattr(first_output, "request_output") or not first_output.request_output:
if not isinstance(first_output, OmniRequestOutput) or not first_output:
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output.request_output
req_out = first_output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
@@ -131,7 +131,7 @@ def _run_inference(
elapsed_ms = (time.time() - start) * 1000
return InferenceResult(
images=outputs[0].request_output.images,
images=outputs[0].images,
elapsed_ms=elapsed_ms,
)
finally:
@@ -49,10 +49,10 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
def _extract_single_image(outputs) -> Image.Image:
first_output = outputs[0]
assert first_output.final_output_type == "image"
if not hasattr(first_output, "request_output") or not first_output.request_output:
if not isinstance(first_output, OmniRequestOutput) or not first_output:
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output.request_output
req_out = first_output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
@@ -1,4 +1,5 @@
import os
from contextlib import contextmanager
from types import SimpleNamespace
import pytest
@@ -10,6 +11,7 @@ from tests.helpers.mark import hardware_test
from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
DistributedAutoencoderKLWan,
select_auto_spatial_shard_split_dim,
)
from vllm_omni.platforms import current_omni_platform
@@ -17,6 +19,23 @@ from vllm_omni.platforms import current_omni_platform
# the end uses hardware_test (H100 x2) plus full_model / diffusion / parallel.
@contextmanager
def _active_spatial_shard(split_dim: str):
token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
wan_spatial_shard.SpatialShardContext(
input_extent=2,
local_input_extent=2,
split_dim=split_dim,
rank=0,
world_size=3,
)
)
try:
yield
finally:
wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
@pytest.mark.core_model
@pytest.mark.cpu
def test_split_for_parallel_decode_pads_uneven_height():
@@ -231,6 +250,128 @@ def test_reshard_from_trimmed_width_pads_invalid_columns(monkeypatch: pytest.Mon
assert torch.equal(out.flatten(), torch.tensor([4.0, 0.0]))
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize("split_dim", ["height", "width"])
@pytest.mark.parametrize(
("full_extent", "world_size"),
[
pytest.param(5, 4, id="four-ranks-empty-tail"),
pytest.param(20, 8, id="production-eight-ranks-empty-tail"),
pytest.param(39, 2, id="two-ranks-partial-tail"),
],
)
def test_reshard_from_trimmed_extent_clamps_and_reconstructs(
monkeypatch: pytest.MonkeyPatch,
split_dim: str,
full_extent: int,
world_size: int,
):
local_extent = (full_extent + world_size - 1) // world_size
shape = (1, 1, 1, full_extent, 1) if split_dim == "height" else (1, 1, 1, 1, full_extent)
x = torch.arange(full_extent, dtype=torch.float32).reshape(shape)
outputs = []
for rank in range(world_size):
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group, rank=rank: (rank, world_size))
token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
wan_spatial_shard.SpatialShardContext(
input_extent=full_extent,
local_input_extent=local_extent,
split_dim=split_dim,
rank=rank,
world_size=world_size,
)
)
try:
out = wan_spatial_shard.reshard_from_trimmed_extent(
x,
local_extent=local_extent,
split_dim=split_dim,
group=object(),
)
finally:
wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
start = rank * local_extent
valid_extent = min(local_extent, max(0, full_extent - start))
expected = torch.zeros(local_extent, dtype=torch.float32)
if valid_extent:
expected[:valid_extent] = torch.arange(start, start + valid_extent, dtype=torch.float32)
assert out.shape[_spatial_dim_for_test(split_dim)] == local_extent
assert torch.equal(out.flatten(), expected)
outputs.append(out)
reconstructed = torch.cat(outputs, dim=_spatial_dim_for_test(split_dim))
reconstructed = reconstructed.narrow(_spatial_dim_for_test(split_dim), 0, full_extent)
assert torch.equal(reconstructed, x)
def _spatial_dim_for_test(split_dim: str) -> int:
return -2 if split_dim == "height" else -1
@pytest.mark.core_model
@pytest.mark.cpu
def test_reshard_from_trimmed_extent_rejects_rank_mismatch(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 2))
x = torch.zeros((1, 1, 1, 1, 4), dtype=torch.float32)
token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
wan_spatial_shard.SpatialShardContext(
input_extent=4,
local_input_extent=2,
split_dim="width",
rank=0,
world_size=2,
)
)
try:
with pytest.raises(RuntimeError, match="group_rank=1, context_rank=0"):
wan_spatial_shard.reshard_from_trimmed_extent(
x,
local_extent=2,
split_dim="width",
group=object(),
)
finally:
wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
@pytest.mark.core_model
@pytest.mark.cpu
def test_attention_wrapper_rejects_spatial_extent_change(monkeypatch: pytest.MonkeyPatch):
class ShrinkingWanAttentionBlock(torch.nn.Module):
def forward(self, x):
return x[..., :-1]
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 2))
monkeypatch.setattr(
wan_spatial_shard,
"all_gather_along_dim",
lambda x, *, group, dim, dst=None: torch.cat([x, x], dim=dim),
)
module = ShrinkingWanAttentionBlock()
wan_spatial_shard._patch_attention_block(module, group=object())
x = torch.zeros((1, 1, 1, 1, 2), dtype=torch.float32)
token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
wan_spatial_shard.SpatialShardContext(
input_extent=4,
local_input_extent=2,
split_dim="width",
rank=0,
world_size=2,
)
)
try:
with pytest.raises(
RuntimeError,
match=r"local_extent=2, gathered_extent=4, trimmed_extent=4, output_extent=3",
):
module(x)
finally:
wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
@pytest.mark.core_model
@pytest.mark.cpu
def test_halo_exchange_single_rank_noop(monkeypatch: pytest.MonkeyPatch):
@@ -255,12 +396,14 @@ def test_dist_zero_pad_only_applies_global_height_edges(monkeypatch: pytest.Monk
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
mid = mid_rank_pad(x)
with _active_spatial_shard("height"):
mid = mid_rank_pad(x)
assert mid.shape == (1, 1, 2, 3)
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object())
last = last_rank_pad(x)
with _active_spatial_shard("height"):
last = last_rank_pad(x)
assert last.shape == (1, 1, 3, 3)
@@ -270,13 +413,15 @@ def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.Monke
x = torch.ones((1, 1, 2, 2))
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
mid = mid_rank_pad(x)
mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
with _active_spatial_shard("width"):
mid = mid_rank_pad(x)
assert mid.shape == (1, 1, 2, 2)
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width")
last = last_rank_pad(x)
last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object())
with _active_spatial_shard("width"):
last = last_rank_pad(x)
assert last.shape == (1, 1, 2, 3)
@@ -299,6 +444,111 @@ def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pyt
assert vae._spatial_shard_decode_enabled(z) is False
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize(
"z_shape,world_size,expected",
[
((1, 16, 1, 60, 104), 4, "width"),
((1, 16, 5, 60, 104), 4, "width"),
((1, 16, 5, 104, 60), 4, "height"),
((1, 16, 5, 2, 2), 4, None),
((1, 16, 5, 64, 64), 4, "width"),
((1, 16, 64, 64), 4, None),
((1, 16, 5, 64, 64), 1, None),
],
)
def test_auto_spatial_shard_selector_is_shape_aware(z_shape, world_size, expected):
assert select_auto_spatial_shard_split_dim(z_shape, world_size) == expected
@pytest.mark.core_model
@pytest.mark.cpu
def test_auto_spatial_shard_gate_requires_full_group(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(
"vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
lambda group=None: 4,
)
vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
vae.use_tiling = True
vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=4, parallel_mode="auto")
vae.is_distributed_enabled = lambda: True
large = torch.zeros((1, 16, 5, 60, 104))
unshardable = torch.zeros((1, 16, 1, 2, 2))
unsupported = torch.zeros((1, 16, 60, 104))
assert vae._spatial_shard_decode_split_dim(large, 4) == "width"
assert vae._spatial_shard_decode_enabled(large) is True
assert vae._spatial_shard_decode_enabled(unshardable) is False
assert vae._spatial_shard_decode_enabled(unsupported) is False
vae.distributed_executor.parallel_size = 2
assert vae._spatial_shard_decode_enabled(large) is False
@pytest.mark.core_model
@pytest.mark.cpu
def test_set_parallel_size_prepares_spatial_wrappers_only_for_a_full_group(monkeypatch: pytest.MonkeyPatch):
vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
executor = SimpleNamespace(group=object(), parallel_size=1, parallel_mode="tile")
def set_parallel_size(parallel_size: int, mode: str = "tile") -> None:
executor.parallel_size = parallel_size
executor.parallel_mode = mode
executor.set_parallel_size = set_parallel_size
object.__setattr__(vae, "distributed_executor", executor)
installed = []
monkeypatch.setattr(
"vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
lambda group=None: 4,
)
monkeypatch.setattr(
wan_spatial_shard,
"install_wan_spatial_shard_decode",
lambda candidate, group, split_dim: installed.append((candidate, group, split_dim)),
)
vae.set_parallel_size(2, mode="auto")
assert installed == []
vae.set_parallel_size(4, mode="auto")
assert installed == [(vae, vae.distributed_executor.group, "height")]
@pytest.mark.core_model
@pytest.mark.cpu
@torch.no_grad()
def test_installed_spatial_decoder_keeps_direct_fallback_and_can_change_axis():
vae = DistributedAutoencoderKLWan(
base_dim=8,
decoder_base_dim=8,
z_dim=4,
dim_mult=[1, 2],
num_res_blocks=1,
attn_scales=[],
temperal_downsample=[False],
latents_mean=[0.0] * 4,
latents_std=[1.0] * 4,
scale_factor_temporal=1,
scale_factor_spatial=2,
).eval()
latents = torch.randn(1, 4, 1, 4, 5)
parameter_names = {name for name, _ in vae.named_parameters()}
state_dict = {name: value.clone() for name, value in vae.state_dict().items()}
reference = vae.decode(latents, return_dict=False)[0]
wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="height")
assert any(getattr(module, "_vllm_omni_dynamic_spatial_shard_conv", False) for module in vae.decoder.modules())
wan_spatial_shard.install_wan_spatial_shard_decode(vae, group=object(), split_dim="width")
assert {name for name, _ in vae.named_parameters()} == parameter_names
vae.load_state_dict(state_dict, strict=True)
assert torch.equal(vae.decode(latents, return_dict=False)[0], reference)
@pytest.mark.core_model
@pytest.mark.cpu
def test_spatial_shard_width_gate_selects_width(monkeypatch: pytest.MonkeyPatch):
@@ -444,3 +694,99 @@ def test_spatial_shard_decode_matches_reference(split_dim: str):
assert mean_abs_diff <= _SPATIAL_SHARD_TOLERANCE, (
f"spatial_shard_{split_dim} mean_abs_diff {mean_abs_diff} exceeds {_SPATIAL_SHARD_TOLERANCE}"
)
def _auto_spatial_shard_decode_worker(rank: int, return_dict, master_port: str) -> None:
from vllm_omni.diffusion.distributed.parallel_state import (
destroy_model_parallel,
init_distributed_environment,
initialize_model_parallel,
)
os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = master_port
device = current_omni_platform.get_torch_device(rank)
current_omni_platform.set_device(device)
backend = current_omni_platform.dist_backend
init_distributed_environment(world_size=2, rank=rank, local_rank=rank, backend=backend)
initialize_model_parallel(sequence_parallel_size=2, ulysses_degree=2, backend=backend)
try:
vae = DistributedAutoencoderKLWan.from_pretrained(
_SPATIAL_SHARD_MODEL,
subfolder=_SPATIAL_SHARD_SUBFOLDER,
torch_dtype=torch.float32,
).to(device=device, dtype=torch.float32)
vae.eval()
generator = torch.Generator(device=device).manual_seed(1)
small = torch.randn((1, vae.config.z_dim, 1, 60, 104), generator=generator, device=device)
landscape = torch.randn((1, vae.config.z_dim, 2, 60, 104), generator=generator, device=device)
portrait = torch.randn((1, vae.config.z_dim, 2, 104, 60), generator=generator, device=device)
with torch.inference_mode():
vae.use_tiling = False
vae.set_parallel_size(1, mode="tile")
small_reference = vae.decode(small, return_dict=False)[0]
vae.use_tiling = True
vae.set_parallel_size(2, mode="tile")
tile_before_install = vae.decode(small, return_dict=False)[0]
vae.set_parallel_size(2, mode="auto")
assert vae._spatial_shard_decode_split_dim(small, 2) == "width"
small_sharded = vae.decode(small, return_dict=False)[0]
assert vae._spatial_shard_decode_split_dim(landscape, 2) == "width"
vae.decode(landscape, return_dict=False)
vae.set_parallel_size(2, mode="tile")
tile_after_install = vae.decode(small, return_dict=False)[0]
vae.use_tiling = False
vae.set_parallel_size(1, mode="tile")
portrait_reference = vae.decode(portrait, return_dict=False)[0]
vae.use_tiling = True
vae.set_parallel_size(2, mode="auto")
assert vae._spatial_shard_decode_split_dim(portrait, 2) == "height"
portrait_sharded = vae.decode(portrait, return_dict=False)[0]
assert all(
buffer is None
for module in vae.decoder.modules()
for name, buffer in module._buffers.items()
if name in {"_halo_recv_top_buf", "_halo_recv_bottom_buf"}
)
if rank == 0:
small_diff = (small_sharded.float() - small_reference.float()).abs()
portrait_diff = (portrait_sharded.float() - portrait_reference.float()).abs()
return_dict["tile_fallback_exact"] = torch.equal(tile_after_install, tile_before_install)
return_dict["small_max_abs_diff"] = small_diff.max().item()
return_dict["small_mean_abs_diff"] = small_diff.mean().item()
return_dict["portrait_max_abs_diff"] = portrait_diff.max().item()
return_dict["portrait_mean_abs_diff"] = portrait_diff.mean().item()
finally:
destroy_model_parallel()
if dist.is_initialized():
dist.destroy_process_group()
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@hardware_test(res={"cuda": "H100"}, num_cards=2)
def test_auto_spatial_shard_switches_shape_paths_safely():
manager = mp.get_context("spawn").Manager()
return_dict = manager.dict()
mp.spawn(
_auto_spatial_shard_decode_worker,
args=(return_dict, "29502"),
nprocs=2,
join=True,
)
assert return_dict.get("tile_fallback_exact") is True
assert return_dict["small_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
assert return_dict["small_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
assert return_dict["portrait_max_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
assert return_dict["portrait_mean_abs_diff"] <= _SPATIAL_SHARD_TOLERANCE
+2 -2
View File
@@ -29,10 +29,10 @@ def test_diffusion_model(tmp_path: Path, omni_runner):
raise ValueError("Empty outputs from Omni.generate()")
first_output = outputs[0]
assert first_output.final_output_type == "image"
if not hasattr(first_output, "request_output") or not first_output.request_output:
if not isinstance(first_output, OmniRequestOutput) or not first_output:
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output.request_output
req_out = first_output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
return req_out.images
@@ -0,0 +1,26 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm_omni.diffusion.layers.norm import RMSNorm
from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLRMSNorm
def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None:
"""Qwen keeps BF16 gamma while native fallback accumulates in FP32."""
eps = 1e-6
norm = MiniMaxH3Qwen3VLRMSNorm(hidden_size=4, eps=eps, dtype=torch.bfloat16)
norm.weight.data.copy_(torch.tensor([1.0, 0.5, 1.5, 2.0], dtype=torch.bfloat16))
x = torch.tensor(
[[1.0, -2.0, 3.0, -4.0], [-5.0, 6.0, -7.0, 8.0]],
dtype=torch.bfloat16,
)
x_fp32 = x.float()
expected = (x_fp32 * torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + eps) * norm.weight.float()).to(x.dtype)
assert isinstance(norm, RMSNorm)
assert norm.weight.dtype == torch.bfloat16
assert set(norm.state_dict()) == {"weight"}
torch.testing.assert_close(norm.forward_native(x), expected, atol=0, rtol=0)
@@ -122,8 +122,8 @@ def test_cpu_offload_diffusion_model(model_name: str):
print(f"No offload peak memory: {no_offload_peak_memory} MB")
if model_name == "stabilityai/stable-audio-open-1.0":
audio_offload = output_offload[0].request_output.multimodal_output.get("audio")
audio_no_offload = output_no_offload[0].request_output.multimodal_output.get("audio")
audio_offload = output_offload[0].multimodal_output.get("audio")
audio_no_offload = output_no_offload[0].multimodal_output.get("audio")
check_audio_determinism(audio_offload, audio_no_offload, atol=1e-2)
# Set platform-specific VRAM saving thresholds to account
@@ -20,18 +20,10 @@ IMAGE_VIDEO_MODELS = {
MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
LAYERWISE_MODEL_CASES = [
pytest.param(
"stabilityai/stable-audio-open-1.0",
marks=[pytest.mark.full_model, pytest.mark.diffusion],
id="stable_audio",
),
pytest.param(
"riverclouds/qwen_image_random",
marks=[pytest.mark.core_model, pytest.mark.diffusion],
id="qwen_image_random",
),
]
MODEL_MARKS = {
"riverclouds/qwen_image_random": pytest.mark.core_model,
"stabilityai/stable-audio-open-1.0": pytest.mark.full_model,
}
AUDIO_MODEL_PARAMS = {
"runner_params": {},
@@ -101,8 +93,9 @@ def run_inference(
return peak, output
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
@pytest.mark.parametrize("model_name", LAYERWISE_MODEL_CASES)
@pytest.mark.parametrize("model_name", list(MODELS.keys()))
def test_layerwise_offload_diffusion_model(model_name: str):
"""Test that layerwise offloading reduces GPU memory usage.
@@ -134,8 +127,8 @@ def test_layerwise_offload_diffusion_model(model_name: str):
print(f"No offload peak memory: {no_offload_peak_memory} MB")
if model_name == "stabilityai/stable-audio-open-1.0":
audio_offload = output_offload[0].request_output.multimodal_output.get("audio")
audio_no_offload = output_no_offload[0].request_output.multimodal_output.get("audio")
audio_offload = output_offload[0].multimodal_output.get("audio")
audio_no_offload = output_no_offload[0].multimodal_output.get("audio")
# Match the sibling cpu-offload test's tolerance: layerwise offload moves
# blocks across the PCIe bus on a side stream, which can perturb cuBLAS
# algorithm selection and produce ~ULP-level drift larger than 1e-3.
@@ -94,8 +94,8 @@ def _generate_single_stage_image(
if hasattr(first_output, "images") and first_output.images:
images = first_output.images
else:
assert hasattr(first_output, "request_output") and first_output.request_output
request_output = first_output.request_output
assert isinstance(first_output, OmniRequestOutput) and first_output
request_output = first_output
if isinstance(request_output, list):
req_out = request_output[0]
else:
@@ -88,8 +88,8 @@ def _generate_single_stage_image(
if hasattr(first_output, "images") and first_output.images:
images = first_output.images
else:
assert hasattr(first_output, "request_output") and first_output.request_output
request_output = first_output.request_output
assert isinstance(first_output, OmniRequestOutput) and first_output
request_output = first_output
if isinstance(request_output, list):
req_out = request_output[0]
else:
@@ -147,10 +147,10 @@ def _generate_single_stage_video(
first = outputs[0]
# Unwrap pipeline-style outputs (multi-stage / OmniRequestOutput.request_output).
# Unwrap pipeline-style outputs (multi-stage / OmniRequestOutput).
frames: Any = None
if hasattr(first, "request_output") and isinstance(first.request_output, list):
inner = first.request_output[0]
if isinstance(first, OmniRequestOutput) and isinstance(first, list):
inner = first[0]
if isinstance(inner, OmniRequestOutput) and inner.images:
frames = inner.images[0]
if frames is None and hasattr(first, "images") and first.images:
@@ -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:
@@ -234,8 +234,8 @@ def _generate_bagel_image(
if images := getattr(req_output, "images", None):
generated_image = images[0]
break
if hasattr(req_output, "request_output") and req_output.request_output:
stage_outputs = req_output.request_output
if isinstance(req_output, OmniRequestOutput) and req_output:
stage_outputs = req_output
if not isinstance(stage_outputs, list):
stage_outputs = [stage_outputs]
for stage_out in stage_outputs:
@@ -250,8 +250,8 @@ def _generate_bagel_image(
# Check LLM stage output — should have finish_reason=stop (not length)
for req_output in omni_outputs:
if hasattr(req_output, "request_output") and req_output.request_output:
stage_outputs = req_output.request_output
if isinstance(req_output, OmniRequestOutput) and req_output:
stage_outputs = req_output
if not isinstance(stage_outputs, list):
stage_outputs = [stage_outputs]
for stage_out in stage_outputs:
@@ -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")
@@ -310,7 +310,7 @@ def _generate_image(omni, config: QualityTestConfig):
first = outputs[0]
if hasattr(first, "images") and first.images:
return first.images[0], peak_mem
inner = first.request_output
inner = first
if inner is not None and hasattr(inner, "images") and inner.images:
return inner.images[0], peak_mem
raise ValueError("Could not extract image from output.")
@@ -341,15 +341,15 @@ def _generate_video(omni, config: QualityTestConfig):
)
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
first = outputs[0]
if hasattr(first, "request_output") and isinstance(first.request_output, list):
inner = first.request_output[0]
if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
frames = inner.images[0] if inner.images else None
else:
frames = inner
elif hasattr(first, "images") and first.images:
frames = first.images[0]
if isinstance(outputs, list) and isinstance(outputs[0], OmniRequestOutput):
first = outputs[0]
elif isinstance(outputs, OmniRequestOutput):
first = outputs
else:
raise ValueError("Could not extract video frames from output.")
if hasattr(first, "images"):
frames = first.images[0] if first.images else None
else:
raise ValueError("Could not extract video frames from output.")
@@ -444,11 +444,11 @@ def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch
image = Image.new("RGB", (2, 2))
inner = OmniRequestOutput.from_diffusion(request_id="req", images=[image])
outer = OmniRequestOutput(
outer = OmniRequestOutput.from_stage_output(
inner,
request_id="req",
stage_id=0,
final_output_type="image",
request_output=inner,
finished=True,
)
@@ -73,11 +73,9 @@ def test_run_summary_reports_worker_peak_memory():
assert summary["per_run_peak_memory_mb"] == [100.0, 150.0]
def test_request_peak_memory_prefers_inner_diffusion_output():
inner = type("InnerOutput", (), {"peak_memory_mb": 321.0})()
outer = type("OuterOutput", (), {"request_output": inner, "peak_memory_mb": 123.0})()
assert _request_peak_memory_mb(outer) == 321.0
def test_request_peak_memory_reads_peak_memory_mb_directly():
result = type("Output", (), {"peak_memory_mb": 123.0})()
assert _request_peak_memory_mb(result) == 123.0
def test_metric_guidance_describes_thresholds():
Loaded 100 of 263 files, more files were not shown because too many files have changed in this diff. Show more