Compare commits
70
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1ae56cfcfb | ||
|
|
14bd8f8772 | ||
|
|
14f2c16fb2 | ||
|
|
9683b7f4b4 | ||
|
|
07ceb3324a | ||
|
|
b4b907e7b6 | ||
|
|
f208f907bc | ||
|
|
284e05c88b | ||
|
|
c3de04e127 | ||
|
|
4cd667875f | ||
|
|
e1c5f59aab | ||
|
|
02069a5bbf | ||
|
|
cd9ae66b3a | ||
|
|
0a1846aa45 | ||
|
|
b283215c28 | ||
|
|
199ef7779c | ||
|
|
3c6648d8a5 | ||
|
|
d7216ef514 | ||
|
|
2fdbf2234a | ||
|
|
265d80b4dd | ||
|
|
5d09cf27a9 | ||
|
|
170d498771 | ||
|
|
2a9f6aaa8e | ||
|
|
5664e69dbf | ||
|
|
e96d3fe2a6 | ||
|
|
309d9c3ea8 | ||
|
|
65b3d41b3a | ||
|
|
4616823125 | ||
|
|
611c083587 | ||
|
|
baba7d1e2e | ||
|
|
9d2bb23ff6 | ||
|
|
3875f8ded4 | ||
|
|
9798bc3a8d | ||
|
|
68bf8c2c42 | ||
|
|
444485650b | ||
|
|
2a315e1a2d | ||
|
|
d1e230c95b | ||
|
|
09b6a1604a | ||
|
|
555336e887 | ||
|
|
136cb27b05 | ||
|
|
ce6ab0ca46 | ||
|
|
f39b8684eb | ||
|
|
e57574ae29 | ||
|
|
7b76b6446d | ||
|
|
a70d4d4b3c | ||
|
|
1260ceaae6 | ||
|
|
6f36fd3b78 | ||
|
|
35cacf020c | ||
|
|
d671e6b0db | ||
|
|
b95c73742d | ||
|
|
fe54ecc388 | ||
|
|
4b6d6d5cf7 | ||
|
|
f7a5146093 | ||
|
|
8ecd1f6d5c | ||
|
|
ffd11c181e | ||
|
|
728640f4d9 | ||
|
|
022d7a7695 | ||
|
|
b300f66656 | ||
|
|
bc73546c06 | ||
|
|
2b08419d8b | ||
|
|
5d50bdc2ab | ||
|
|
8a5bf2fb3f | ||
|
|
596c16a550 | ||
|
|
90a08c0eda | ||
|
|
dbc0dd6d85 | ||
|
|
83d1f11449 | ||
|
|
4b66f62e10 | ||
|
|
f41451f64c | ||
|
|
bbe6ccc512 | ||
|
|
bbb709cfb5 |
No files matched your search
@@ -1,153 +1,103 @@
|
||||
#!/bin/bash
|
||||
# vllm-omni customized version
|
||||
# Based on: vllm/.buildkite/scripts/hardware_ci/run-amd-test.sh
|
||||
# Last synced: 2025-12-15
|
||||
# Modifications: docker image name for vllm-omni
|
||||
|
||||
# This script runs test inside the corresponding ROCm docker container.
|
||||
set -o pipefail
|
||||
# Run vLLM-Omni ROCm tests directly in the MI300 Kubernetes pod. The pod's
|
||||
# container image is selected by test-template-amd-omni.j2; MI300 has no DinD.
|
||||
set -euo pipefail
|
||||
|
||||
# Export Python path
|
||||
export PYTHONPATH=".."
|
||||
: "${PYTHONFAULTHANDLER:=1}"
|
||||
: "${HF_HOME:=/home/buildkite-agent/huggingface}"
|
||||
: "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
|
||||
: "${HF_HUB_ETAG_TIMEOUT:=60}"
|
||||
: "${MIOPEN_DEBUG_CONV_DIRECT:=0}"
|
||||
: "${MIOPEN_DEBUG_CONV_GEMM:=0}"
|
||||
: "${VLLM_ROCM_USE_AITER:=0}"
|
||||
export PYTHONFAULTHANDLER HF_HOME HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
|
||||
export MIOPEN_DEBUG_CONV_DIRECT MIOPEN_DEBUG_CONV_GEMM VLLM_ROCM_USE_AITER
|
||||
export PYTORCH_ROCM_ARCH=""
|
||||
export PYTHONPATH="${PYTHONPATH:-..}"
|
||||
|
||||
if [[ "${VLLM_CI_DOCKER_DISABLED:-0}" != "1" ]]; then
|
||||
echo "Error: MI300 CI must run natively with Docker disabled." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [[ -n "${TEST_COMMANDS:-}" ]]; then
|
||||
commands="${TEST_COMMANDS}"
|
||||
else
|
||||
commands="$*"
|
||||
fi
|
||||
if [[ -z "${commands}" ]]; then
|
||||
echo "Error: No test commands provided." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "--- Native in-pod ROCm CI"
|
||||
|
||||
job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-local}}"
|
||||
job_id="${job_id//[^A-Za-z0-9_.-]/_}"
|
||||
native_cache_root="/tmp/vllm-omni-native-${job_id}"
|
||||
TMPDIR="${native_cache_root}/tmp"
|
||||
TORCHINDUCTOR_CACHE_DIR="${native_cache_root}/torchinductor"
|
||||
TRITON_CACHE_DIR="${native_cache_root}/triton"
|
||||
VLLM_CACHE_ROOT="${native_cache_root}/vllm"
|
||||
XDG_CACHE_HOME="${native_cache_root}/xdg"
|
||||
HF_DATASETS_CACHE="${native_cache_root}/huggingface/datasets"
|
||||
export TMPDIR TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT
|
||||
export XDG_CACHE_HOME HF_DATASETS_CACHE
|
||||
mkdir -p "${TMPDIR}" "${TORCHINDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}" \
|
||||
"${VLLM_CACHE_ROOT}" "${XDG_CACHE_HOME}" "${HF_HOME}" \
|
||||
"${HF_DATASETS_CACHE}"
|
||||
|
||||
if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then
|
||||
if ! command -v findmnt >/dev/null 2>&1; then
|
||||
echo "Error: findmnt is required to verify the Hugging Face cache mount." >&2
|
||||
exit 1
|
||||
fi
|
||||
hf_mount=$(findmnt -n -T "${HF_HOME}" -o TARGET 2>/dev/null || true)
|
||||
if [[ -z "${hf_mount}" || "${hf_mount}" == "/" ]]; then
|
||||
echo "Error: MI300 CI requires a persistent volume at or above ${HF_HOME}." >&2
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
# Print ROCm version
|
||||
echo "--- ROCm info"
|
||||
rocminfo
|
||||
|
||||
# cleanup older docker images
|
||||
cleanup_docker() {
|
||||
# Get Docker's root directory
|
||||
docker_root=$(docker info -f '{{.DockerRootDir}}')
|
||||
if [ -z "$docker_root" ]; then
|
||||
echo "Failed to determine Docker root directory."
|
||||
exit 1
|
||||
fi
|
||||
echo "Docker root directory: $docker_root"
|
||||
# Check disk usage of the filesystem where Docker's root directory is located
|
||||
disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
|
||||
# Define the threshold
|
||||
threshold=70
|
||||
if [ "$disk_usage" -gt "$threshold" ]; then
|
||||
echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
|
||||
# Remove dangling images (those that are not tagged and not used by any container)
|
||||
docker image prune -f
|
||||
# Remove unused volumes / force the system prune for old images as well.
|
||||
docker volume prune -f && docker system prune --force --filter "until=72h" --all
|
||||
echo "Docker images and volumes cleanup completed."
|
||||
else
|
||||
echo "Disk usage is below $threshold%. No cleanup needed."
|
||||
fi
|
||||
}
|
||||
expected_gpus="${VLLM_CI_EXPECTED_GPU_COUNT:-1}"
|
||||
python3 - <<'PY'
|
||||
import os
|
||||
|
||||
# Call the cleanup docker function
|
||||
cleanup_docker
|
||||
import torch
|
||||
|
||||
echo "--- Pulling container"
|
||||
## Temporary change to use AMD Docker Hub to store the vllm-omni image
|
||||
# to bypass the rate limit issue with ECR Public Gallery.
|
||||
# Images are now stored in a separate repository for vllm-omni, instead of vllm-ci.
|
||||
# TODO: @tjtanaa point back to ECR Public Gallery
|
||||
# once the amd agents are configured to use ECR Public Gallery.
|
||||
# image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
|
||||
image_name="rocm/vllm-omni:${BUILDKITE_COMMIT}"
|
||||
container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
|
||||
expected = int(os.environ.get("VLLM_CI_EXPECTED_GPU_COUNT", "1"))
|
||||
assert torch.version.hip, "PyTorch is not a ROCm build"
|
||||
assert torch.cuda.is_available(), "ROCm GPU is not available to PyTorch"
|
||||
actual = torch.cuda.device_count()
|
||||
assert actual == expected, f"Expected {expected} ROCm GPU(s), found {actual}"
|
||||
PY
|
||||
|
||||
# TODO: @tjtanaa uncomment this once the amd agents are configured to use ECR Public Gallery.
|
||||
# # Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
|
||||
# sudo apt-get update && sudo apt-get install -y awscli
|
||||
## Use safe docker login helper to prevent race conditions
|
||||
# source "$(dirname "${BASH_SOURCE[0]}")/../../common/scripts/docker_login_ecr_public.sh"
|
||||
# safe_docker_login_ecr_public
|
||||
|
||||
## Pull the container from AMD Docker Hub
|
||||
|
||||
docker pull "${image_name}"
|
||||
|
||||
remove_docker_container() {
|
||||
docker ps -aq --filter "name=${container_name}" | xargs -r docker rm -f || true
|
||||
docker image rm -f "${image_name}" || true
|
||||
}
|
||||
trap remove_docker_container EXIT
|
||||
|
||||
echo "--- Running container"
|
||||
|
||||
HF_CACHE="$(realpath ~)/huggingface"
|
||||
mkdir -p "${HF_CACHE}"
|
||||
HF_MOUNT="/root/.cache/huggingface"
|
||||
|
||||
if [[ -n "${TEST_COMMAND:-}" ]]; then
|
||||
commands="$TEST_COMMAND"
|
||||
else
|
||||
commands="$@"
|
||||
fi
|
||||
echo "Commands:$commands"
|
||||
|
||||
PARALLEL_JOB_COUNT=8
|
||||
MYPYTHONPATH=".."
|
||||
|
||||
# Test that we're launching on the machine that has
|
||||
# proper access to GPUs
|
||||
render_gid=$(getent group render | cut -d: -f3)
|
||||
if [[ -z "$render_gid" ]]; then
|
||||
echo "Error: 'render' group not found. This is required for GPU access." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs.
|
||||
# TODO: @tjtanaa reenable to run VLLM_ROCM_USE_AITER=1 when AITER is shipped with prebuilt kernels.
|
||||
if [[ $commands == *"--shard-id="* ]]; then
|
||||
# assign job count as the number of shards used
|
||||
commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g')
|
||||
for GPU in $(seq 0 $(($PARALLEL_JOB_COUNT-1))); do
|
||||
# assign shard-id for each shard
|
||||
commands_gpu=$(echo "$commands" | sed -E "s/--shard-id[[:blank:]]*=[[:blank:]]*[0-9]*/--shard-id=${GPU} /g" | sed 's/ \\ / /g')
|
||||
echo "Shard ${GPU} commands:$commands_gpu"
|
||||
echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
|
||||
docker run \
|
||||
--device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
|
||||
--network=host \
|
||||
--shm-size=16gb \
|
||||
--group-add "$render_gid" \
|
||||
-e MIOPEN_DEBUG_CONV_DIRECT=0 \
|
||||
-e MIOPEN_DEBUG_CONV_GEMM=0 \
|
||||
-e VLLM_ROCM_USE_AITER=0 \
|
||||
-e HIP_VISIBLE_DEVICES="${GPU}" \
|
||||
-e HF_TOKEN \
|
||||
-v "${HF_CACHE}:${HF_MOUNT}" \
|
||||
-e "HF_HOME=${HF_MOUNT}" \
|
||||
-e "PYTHONPATH=${MYPYTHONPATH}" \
|
||||
--name "${container_name}_${GPU}" \
|
||||
"${image_name}" \
|
||||
/bin/bash -c "${commands_gpu}" \
|
||||
|& while read -r line; do echo ">>Shard $GPU: $line"; done &
|
||||
PIDS+=($!)
|
||||
done
|
||||
#wait for all processes to finish and collect exit codes
|
||||
for pid in "${PIDS[@]}"; do
|
||||
wait "${pid}"
|
||||
STATUS+=($?)
|
||||
done
|
||||
for st in "${STATUS[@]}"; do
|
||||
if [[ ${st} -ne 0 ]]; then
|
||||
echo "One of the processes failed with $st"
|
||||
exit "${st}"
|
||||
echo "Commands:${commands}"
|
||||
if /bin/bash -o pipefail -c '
|
||||
set -E
|
||||
test_status=0
|
||||
trap '\''
|
||||
command_status=$?
|
||||
# Keep running subsequent commands, but retain a failure for the final exit.
|
||||
# Prefer a test failure over pytest "no tests collected" when both occur.
|
||||
if (( test_status == 0 || (test_status == 5 && command_status != 5) )); then
|
||||
test_status=${command_status}
|
||||
fi
|
||||
done
|
||||
'\'' ERR
|
||||
eval "$1"
|
||||
exit "${test_status}"
|
||||
' _ "${commands}"; then
|
||||
exit 0
|
||||
else
|
||||
echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
|
||||
docker run \
|
||||
--device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
|
||||
--network=host \
|
||||
--shm-size=16gb \
|
||||
--group-add "$render_gid" \
|
||||
-e MIOPEN_DEBUG_CONV_DIRECT=0 \
|
||||
-e MIOPEN_DEBUG_CONV_GEMM=0 \
|
||||
-e VLLM_ROCM_USE_AITER=0 \
|
||||
-e HF_TOKEN \
|
||||
-v "${HF_CACHE}:${HF_MOUNT}" \
|
||||
-e "HF_HOME=${HF_MOUNT}" \
|
||||
-e "PYTHONPATH=${MYPYTHONPATH}" \
|
||||
--name "${container_name}" \
|
||||
"${image_name}" \
|
||||
/bin/bash -c "${commands}"
|
||||
exit_code=$?
|
||||
fi
|
||||
|
||||
if [[ ${exit_code} -eq 5 && "${VLLM_CI_ALLOW_NO_TESTS:-0}" == "1" ]]; then
|
||||
echo "Pytest collected no tests; VLLM_CI_ALLOW_NO_TESTS=1, treating exit code 5 as success."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
exit "${exit_code}"
|
||||
@@ -7,54 +7,64 @@ steps:
|
||||
|
||||
- group: ":card_index_dividers: Simple Test"
|
||||
steps:
|
||||
- label: "Simple · Diffusion & Model Executor Test (1/2)"
|
||||
agent_pool: mi325_1
|
||||
- label: "Simple · Model Executor Test · Shard %N/%t"
|
||||
agent_pool: mi300_1
|
||||
parallelism: 2
|
||||
depends_on: amd-build
|
||||
timeout_in_minutes: 70
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
# CPU tests on ROCm gpu machines are very slow, so we give them more time.
|
||||
- "timeout 60m pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- "pytest -sv tests/model_executor -m 'core_model and cpu' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:coverage-model-executor-$$BUILDKITE_PARALLEL_JOB.xml"
|
||||
|
||||
- label: "Simple · Diffusion & Model Executor Test (2/2)"
|
||||
agent_pool: mi325_1
|
||||
- label: "Simple · Diffusion Test · Shard %N/%t"
|
||||
agent_pool: mi300_1
|
||||
parallelism: 2
|
||||
depends_on: amd-build
|
||||
timeout_in_minutes: 70
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# CPU tests on ROCm gpu machines are very slow, so we give them more time.
|
||||
- "timeout 60m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:coverage-diffusion-$$BUILDKITE_PARALLEL_JOB.xml"
|
||||
|
||||
- label: "Simple · Engine&Entrypoints Test"
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
- "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Simple · Other Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- group: ":card_index_dividers: Diffusion Test"
|
||||
steps:
|
||||
- label: "Diffusion · LoRA Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
artifact_paths:
|
||||
- lora-server.log
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
|
||||
- pytest -sv tests/diffusion/lora/ -m 'advanced_model and cuda' --run-level "advanced_model" 2>&1 | tee lora-server.log
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -62,7 +72,7 @@ steps:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
@@ -72,7 +82,7 @@ steps:
|
||||
'
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
|
||||
agent_pool: mi325_4
|
||||
agent_pool: mi300_4
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -84,7 +94,7 @@ steps:
|
||||
|
||||
# merge-only tests
|
||||
- label: "Diffusion Tensor Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
agent_pool: mi300_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -92,7 +102,7 @@ steps:
|
||||
- pytest -s -v tests/diffusion/distributed/test_tensor_parallel.py
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -100,7 +110,7 @@ steps:
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
agent_pool: mi300_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -108,7 +118,7 @@ steps:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -121,7 +131,7 @@ steps:
|
||||
|
||||
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
|
||||
# - label: "Audio Generation Model Test"
|
||||
# agent_pool: mi325_1
|
||||
# agent_pool: mi300_1
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
@@ -132,7 +142,7 @@ steps:
|
||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
|
||||
|
||||
- label: "Engine Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -147,7 +157,7 @@ steps:
|
||||
- group: ":card_index_dividers: E2E Test"
|
||||
steps:
|
||||
- label: "Omni Model Test Qwen2-5-Omni"
|
||||
agent_pool: mi325_4
|
||||
agent_pool: mi300_4
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
commands:
|
||||
@@ -155,7 +165,7 @@ steps:
|
||||
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
|
||||
|
||||
- label: "Omni Model Test Qwen3-Omni"
|
||||
agent_pool: mi325_2
|
||||
agent_pool: mi300_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -163,22 +173,23 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
|
||||
- label: "Qwen3-TTS CustomVoice E2E Test"
|
||||
agent_pool: mi325_1
|
||||
- label: "Qwen3-TTS CustomVoice E2E Test · Shard %N/%t"
|
||||
agent_pool: mi300_1
|
||||
parallelism: 2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 50m bash -c '
|
||||
bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
'
|
||||
|
||||
- label: "Qwen3-TTS Base E2E Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -196,7 +207,7 @@ steps:
|
||||
# Need to debug before reneable numerical changes across large PRs
|
||||
# # split Bagel Model Test with H100 (Real Weights) into three tests
|
||||
# - label: "Bagel Text2Img Model Test (1/3)"
|
||||
# agent_pool: mi325_1
|
||||
# agent_pool: mi300_1
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
@@ -208,7 +219,7 @@ steps:
|
||||
# - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
|
||||
|
||||
# - label: "Bagel Img2Img Model Test (2/3)"
|
||||
# agent_pool: mi325_1
|
||||
# agent_pool: mi300_1
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
@@ -220,7 +231,7 @@ steps:
|
||||
# - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
|
||||
|
||||
# - label: "Bagel Online Serving Test (3/3)"
|
||||
# agent_pool: mi325_1
|
||||
# agent_pool: mi300_1
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
|
||||
@@ -7,8 +7,9 @@ steps:
|
||||
|
||||
- group: ":card_index_dividers: Simple Test"
|
||||
steps:
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
agent_pool: mi325_1
|
||||
- label: "Simple · Model Executor Test · Shard %N/%t"
|
||||
agent_pool: mi300_1
|
||||
parallelism: 2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -16,10 +17,30 @@ steps:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- "pytest -sv tests/model_executor -m 'core_model and cpu' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:coverage-model-executor-$$BUILDKITE_PARALLEL_JOB.xml"
|
||||
|
||||
- label: "Simple · Diffusion Test"
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Simple · Engine&Entrypoints Test"
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
- "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Simple · Other Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -27,19 +48,19 @@ steps:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Custom Pipeline Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model"
|
||||
- pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model"
|
||||
|
||||
|
||||
- label: "Voxtral TTS CUDA Unit Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -49,33 +70,25 @@ steps:
|
||||
- group: ":card_index_dividers: Diffusion Test"
|
||||
steps:
|
||||
- label: "Diffusion · Batch Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
|
||||
- pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
|
||||
|
||||
- label: "Diffusion · Offloader Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
- pytest -s -v tests/diffusion/offloader -m "core_model and cuda" --run-level "core_model"
|
||||
|
||||
- label: "Diffusion · Cache Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -83,10 +96,10 @@ steps:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
|
||||
- pytest -sv tests/diffusion/cache -m 'core_model and cuda' --run-level "core_model"
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
agent_pool: mi300_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -96,7 +109,7 @@ steps:
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
agent_pool: mi300_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -104,7 +117,7 @@ steps:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Diffusion · Model Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -112,7 +125,7 @@ steps:
|
||||
- timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -126,7 +139,7 @@ steps:
|
||||
|
||||
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
|
||||
# - label: "Audio Generation Model Test"
|
||||
# agent_pool: mi325_1
|
||||
# agent_pool: mi300_1
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
@@ -137,7 +150,7 @@ steps:
|
||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
|
||||
|
||||
- label: "Engine Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -149,7 +162,7 @@ steps:
|
||||
'
|
||||
|
||||
# - label: "Omni Model Test Qwen3-Omni"
|
||||
# agent_pool: mi325_2
|
||||
# agent_pool: mi300_2
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
@@ -162,7 +175,7 @@ steps:
|
||||
- group: ":card_index_dividers: E2E Test"
|
||||
steps:
|
||||
- label: "Qwen3-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -173,7 +186,7 @@ steps:
|
||||
- timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
|
||||
|
||||
- label: "CosyVoice3-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
@@ -186,7 +199,7 @@ steps:
|
||||
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
|
||||
# Need to debug before reneable numerical changes across large PRs
|
||||
# - label: "Bagel Text2Img Model Test"
|
||||
# agent_pool: mi325_1
|
||||
# agent_pool: mi300_1
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
@@ -198,7 +211,7 @@ steps:
|
||||
# - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
|
||||
# - label: "Bagel Img2Img Model Test"
|
||||
# agent_pool: mi325_1
|
||||
# agent_pool: mi300_1
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
@@ -210,7 +223,7 @@ steps:
|
||||
# - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
|
||||
# - label: "Bagel Online Serving Test"
|
||||
# agent_pool: mi325_1
|
||||
# agent_pool: mi300_1
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
|
||||
@@ -1,11 +1,65 @@
|
||||
{# vllm-omni customized version
|
||||
Based on: https://github.com/vllm-project/ci-infra/blob/main/buildkite/test-template-amd.j2
|
||||
Last synced: 2025-12-15
|
||||
Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests
|
||||
Last synced: 2026-08-04
|
||||
Modifications: Use the shared MI300 queues and Omni command-passing contract
|
||||
#}
|
||||
{% set docker_image_amd = "rocm/vllm-omni:$BUILDKITE_COMMIT" %}
|
||||
{% set default_working_dir = "/app/vllm-omni" %}
|
||||
|
||||
{# MI300 runs tests natively in Kubernetes; Docker-in-Docker is unavailable. #}
|
||||
{% macro amd_native_k8s_plugin(indent, step) -%}
|
||||
{{ indent }}plugins:
|
||||
{{ indent }} - kubernetes:
|
||||
{{ indent }} podSpecPatch:
|
||||
{{ indent }} automountServiceAccountToken: false
|
||||
{{ indent }} securityContext:
|
||||
{{ indent }} seccompProfile:
|
||||
{{ indent }} type: RuntimeDefault
|
||||
{{ indent }} imagePullSecrets:
|
||||
{{ indent }} - name: docker-config
|
||||
{{ indent }} containers:
|
||||
{{ indent }} - name: container-0
|
||||
{{ indent }} image: {{ docker_image_amd }}
|
||||
{{ indent }} imagePullPolicy: Always
|
||||
{{ indent }} securityContext:
|
||||
{{ indent }} allowPrivilegeEscalation: false
|
||||
{{ indent }} capabilities:
|
||||
{{ indent }} add:
|
||||
{{ indent }} - IPC_LOCK
|
||||
{{ indent }} resources:
|
||||
{{ indent }} limits:
|
||||
{{ indent }} amd.com/gpu: "{{ step.agent_pool[-1:] }}"
|
||||
{{ indent }} requests:
|
||||
{{ indent }} amd.com/gpu: "{{ step.agent_pool[-1:] }}"
|
||||
{{ indent }} volumeMounts:
|
||||
{{ indent }} - name: devshm
|
||||
{{ indent }} mountPath: /dev/shm
|
||||
{{ indent }} env:
|
||||
{{ indent }} - name: AMD_CI_RUNTIME
|
||||
{{ indent }} value: native
|
||||
{{ indent }} - name: NATIVE_CI
|
||||
{{ indent }} value: "true"
|
||||
{{ indent }} - name: VLLM_CI_DOCKER_DISABLED
|
||||
{{ indent }} value: "1"
|
||||
{{ indent }} - name: VLLM_CI_EXPECTED_GPU_COUNT
|
||||
{{ indent }} value: "{{ step.agent_pool[-1:] }}"
|
||||
{{ indent }} - name: HF_HOME
|
||||
{{ indent }} value: /home/buildkite-agent/huggingface
|
||||
{{ indent }} - name: MIOPEN_DEBUG_CONV_DIRECT
|
||||
{{ indent }} value: "0"
|
||||
{{ indent }} - name: MIOPEN_DEBUG_CONV_GEMM
|
||||
{{ indent }} value: "0"
|
||||
{{ indent }} - name: VLLM_ROCM_USE_AITER
|
||||
{{ indent }} value: "0"
|
||||
{{ indent }} - name: PYTORCH_ROCM_ARCH
|
||||
{{ indent }} value: ""
|
||||
{{ indent }} volumes:
|
||||
{{ indent }} - name: devshm
|
||||
{{ indent }} emptyDir:
|
||||
{{ indent }} medium: Memory
|
||||
{{ indent }} sizeLimit: 16Gi
|
||||
{%- endmacro %}
|
||||
|
||||
{% if env %}
|
||||
env:
|
||||
{% for name in env %}
|
||||
@@ -53,19 +107,31 @@ steps:
|
||||
{% if step.agent_pool %}
|
||||
queue: amd_{{ step.agent_pool }}
|
||||
{% else %}
|
||||
queue: amd_mi325_1
|
||||
queue: amd_mi300_1
|
||||
{% endif %}
|
||||
{{ amd_native_k8s_plugin(" ", step) }}
|
||||
{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
|
||||
{% set indented_cmd = cmd_body | replace("\n", "\n ") %}
|
||||
command: bash .buildkite/amd/scripts/run-amd-test.sh
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
AMD_CI_RUNTIME: "native"
|
||||
NATIVE_CI: "true"
|
||||
VLLM_CI_DOCKER_DISABLED: "1"
|
||||
VLLM_CI_EXPECTED_GPU_COUNT: "{{ step.agent_pool[-1:] }}"
|
||||
VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE: "1"
|
||||
HF_HOME: "/home/buildkite-agent/huggingface"
|
||||
HF_HUB_DOWNLOAD_TIMEOUT: "300"
|
||||
HF_HUB_ETAG_TIMEOUT: "60"
|
||||
MIOPEN_DEBUG_CONV_DIRECT: "0"
|
||||
MIOPEN_DEBUG_CONV_GEMM: "0"
|
||||
VLLM_ROCM_USE_AITER: "0"
|
||||
PYTORCH_ROCM_ARCH: ""
|
||||
{% if step.env %}
|
||||
{% for name in step.env %}
|
||||
{{ name }}: "{{ step.env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
TEST_COMMAND: |-
|
||||
TEST_COMMANDS: |-
|
||||
(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
|
||||
{% if env %}
|
||||
{% for name in env %}
|
||||
@@ -82,6 +148,15 @@ steps:
|
||||
{% endif %}
|
||||
{{ indented_cmd | safe }}
|
||||
priority: 100
|
||||
{% if step.parallelism %}
|
||||
parallelism: {{ step.parallelism }}
|
||||
{% endif %}
|
||||
{% if step.artifact_paths %}
|
||||
artifact_paths:
|
||||
{% for path in step.artifact_paths %}
|
||||
- "{{ path }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
{% if step.timeout_in_minutes %}
|
||||
timeout_in_minutes: {{ step.timeout_in_minutes }}
|
||||
{% endif %}
|
||||
@@ -100,19 +175,31 @@ steps:
|
||||
{% if step.agent_pool %}
|
||||
queue: amd_{{ step.agent_pool }}
|
||||
{% else %}
|
||||
queue: amd_mi325_1
|
||||
queue: amd_mi300_1
|
||||
{% endif %}
|
||||
{{ amd_native_k8s_plugin(" ", step) }}
|
||||
{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
|
||||
{% set indented_cmd = cmd_body | replace("\n", "\n ") %}
|
||||
command: bash .buildkite/amd/scripts/run-amd-test.sh
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
AMD_CI_RUNTIME: "native"
|
||||
NATIVE_CI: "true"
|
||||
VLLM_CI_DOCKER_DISABLED: "1"
|
||||
VLLM_CI_EXPECTED_GPU_COUNT: "{{ step.agent_pool[-1:] }}"
|
||||
VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE: "1"
|
||||
HF_HOME: "/home/buildkite-agent/huggingface"
|
||||
HF_HUB_DOWNLOAD_TIMEOUT: "300"
|
||||
HF_HUB_ETAG_TIMEOUT: "60"
|
||||
MIOPEN_DEBUG_CONV_DIRECT: "0"
|
||||
MIOPEN_DEBUG_CONV_GEMM: "0"
|
||||
VLLM_ROCM_USE_AITER: "0"
|
||||
PYTORCH_ROCM_ARCH: ""
|
||||
{% if step.env %}
|
||||
{% for name in step.env %}
|
||||
{{ name }}: "{{ step.env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
TEST_COMMAND: |-
|
||||
TEST_COMMANDS: |-
|
||||
(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
|
||||
{% if env %}
|
||||
{% for name in env %}
|
||||
@@ -129,6 +216,15 @@ steps:
|
||||
{% endif %}
|
||||
{{ indented_cmd | safe }}
|
||||
priority: 100
|
||||
{% if step.parallelism %}
|
||||
parallelism: {{ step.parallelism }}
|
||||
{% endif %}
|
||||
{% if step.artifact_paths %}
|
||||
artifact_paths:
|
||||
{% for path in step.artifact_paths %}
|
||||
- "{{ path }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
{% if step.timeout_in_minutes %}
|
||||
timeout_in_minutes: {{ step.timeout_in_minutes }}
|
||||
{% endif %}
|
||||
|
||||
@@ -9,6 +9,11 @@
|
||||
# a2b3_npu_1 — ascend-a2b3, resource_class npu-1 (A2/B3 CI image)
|
||||
# a2b3_npu_4 — ascend-a2b3, resource_class npu-4
|
||||
# a3_npu_2 — ascend-a3, resource_class npu-2
|
||||
#
|
||||
# The ``l4_*`` presets run under the ``docker`` plugin, which does not put
|
||||
# ``buildkite-agent`` on PATH inside the container the way the ``kubernetes``
|
||||
# presets do, so they set ``mount-buildkite-agent: true``; without it any step
|
||||
# calling ``buildkite-agent`` (artifact upload, annotate, meta-data) fails there.
|
||||
|
||||
x-npu-k8s-plugins: &npu_k8s_plugins
|
||||
- kubernetes:
|
||||
@@ -26,6 +31,7 @@ mirror_hardwares:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
mount-buildkite-agent: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- HF_HOME=/fsx/hf_cache
|
||||
@@ -41,6 +47,7 @@ mirror_hardwares:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
mount-buildkite-agent: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- HF_HOME=/fsx/hf_cache
|
||||
|
||||
Executable
+82
@@ -0,0 +1,82 @@
|
||||
#!/usr/bin/env bash
|
||||
# Run a model's E2E tests once per entry mode as separate coverage runs and upload
|
||||
# the reports, so per-mode attribution survives (see vllm-project/vllm-omni#5332).
|
||||
#
|
||||
# Usage: run_cov_split.sh --model-id <id> [--offline <paths>] [--online <paths>] \
|
||||
# -- <pytest args...>
|
||||
#
|
||||
# --model-id Names the artifacts: coverage-<id>-{offline,online}-<step id>.xml.gz.
|
||||
# Use the model's directory name under vllm_omni/*/models/.
|
||||
# --offline Pytest path(s) for the offline half. Repeatable.
|
||||
# --online Pytest path(s) for the online half. Repeatable.
|
||||
# -- Everything after this is forwarded verbatim to every pytest run,
|
||||
# quoting preserved: the job's own -m expression, --run-level, and
|
||||
# anything else it needs (e.g. --test-config-file).
|
||||
#
|
||||
# At least one of --offline/--online is required, so a job with only one mode runs
|
||||
# just that half. Total runtime is bounded by the step's timeout_in_minutes.
|
||||
#
|
||||
# Exits non-zero if any half or the upload failed.
|
||||
set -uo pipefail
|
||||
|
||||
MODEL_ID=""
|
||||
OFFLINE=()
|
||||
ONLINE=()
|
||||
PYTEST_ARGS=()
|
||||
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
--model-id) MODEL_ID="$2"; shift 2 ;;
|
||||
--offline) OFFLINE+=("$2"); shift 2 ;;
|
||||
--online) ONLINE+=("$2"); shift 2 ;;
|
||||
--) shift; PYTEST_ARGS=("$@"); break ;;
|
||||
*) echo "run_cov_split.sh: unknown argument: $1" >&2; exit 2 ;;
|
||||
esac
|
||||
done
|
||||
|
||||
if [[ -z "${MODEL_ID}" ]]; then
|
||||
echo "run_cov_split.sh: missing --model-id" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
if ((${#OFFLINE[@]} == 0 && ${#ONLINE[@]} == 0)); then
|
||||
echo "run_cov_split.sh: need at least one of --offline / --online" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
# Two steps can test the same model, and artifacts are scoped to the build (a
|
||||
# nightly one carries the ready, merge and nightly tiers at once), so unsuffixed
|
||||
# reports would share a path -- a glob download then races to rename onto one file.
|
||||
# Not BUILDKITE_JOB_ID: the step id survives retries, so attempts collapse under
|
||||
# Buildkite's latest-attempt lookup instead of accumulating. A parallel step would
|
||||
# need BUILDKITE_PARALLEL_JOB as well; no pilot is one.
|
||||
REPORT_SUFFIX="${BUILDKITE_STEP_ID:-local}"
|
||||
OFFLINE_REPORT="coverage-${MODEL_ID}-offline-${REPORT_SUFFIX}.xml"
|
||||
ONLINE_REPORT="coverage-${MODEL_ID}-online-${REPORT_SUFFIX}.xml"
|
||||
|
||||
run_half() {
|
||||
local mode="$1"
|
||||
shift
|
||||
local report="coverage-${MODEL_ID}-${mode}-${REPORT_SUFFIX}.xml"
|
||||
echo "--- coverage: ${MODEL_ID} ${mode}"
|
||||
pytest -s -v "$@" "${PYTEST_ARGS[@]}" --cov=vllm_omni --cov-report="xml:${report}"
|
||||
local status=$?
|
||||
# Cobertura lists every statement in vllm_omni whether or not it ran, so the
|
||||
# report is ~7MB no matter how little the job covers. The repetition is what
|
||||
# dominates it, and gzip takes that to ~400KB.
|
||||
gzip -9 -f "${report}" || status=1
|
||||
return "${status}"
|
||||
}
|
||||
|
||||
# Clear both modes, not just the ones being run: the upload glob matches both, so
|
||||
# a report left by an earlier run would otherwise be published as if the mode it
|
||||
# belongs to had run this time.
|
||||
rm -f "${OFFLINE_REPORT}" "${ONLINE_REPORT}" "${OFFLINE_REPORT}.gz" "${ONLINE_REPORT}.gz"
|
||||
|
||||
EXIT=0
|
||||
((${#OFFLINE[@]})) && { run_half offline "${OFFLINE[@]}" || EXIT=1; }
|
||||
((${#ONLINE[@]})) && { run_half online "${ONLINE[@]}" || EXIT=1; }
|
||||
|
||||
buildkite-agent artifact upload "coverage-${MODEL_ID}-*-${REPORT_SUFFIX}.xml.gz" || EXIT=1
|
||||
|
||||
exit "${EXIT}"
|
||||
@@ -76,6 +76,9 @@ steps:
|
||||
- tests/e2e/online_serving/test_qwen3_tts_customvoice.py
|
||||
- tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
|
||||
- vllm_omni/model_executor/models/qwen3_tts/
|
||||
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
|
||||
- vllm_omni/model_executor/models/common/snake_activation.py
|
||||
- vllm_omni/model_executor/models/common/alias_free_activation.py
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
|
||||
- vllm_omni/deploy/qwen3_tts.yaml
|
||||
commands:
|
||||
@@ -92,31 +95,24 @@ steps:
|
||||
- tests/e2e/online_serving/test_qwen3_tts_base.py
|
||||
- tests/e2e/offline_inference/test_qwen3_tts_base.py
|
||||
- vllm_omni/model_executor/models/qwen3_tts/
|
||||
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
|
||||
- vllm_omni/model_executor/models/common/snake_activation.py
|
||||
- vllm_omni/model_executor/models/common/alias_free_activation.py
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
|
||||
- vllm_omni/deploy/qwen3_tts.yaml
|
||||
- .buildkite/common/scripts/run_cov_split.sh
|
||||
- pyproject.toml
|
||||
# 40 -> 60: the coverage split runs the model twice.
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "TTS · Higgs-Audio-V3 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_higgs_audio_v3.py
|
||||
- vllm_omni/model_executor/models/higgs_audio_v3/
|
||||
- vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py
|
||||
- vllm_omni/deploy/higgs_multimodal_qwen3.yaml
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_USE_DEEP_GEMM=0
|
||||
export VLLM_MOE_USE_DEEP_GEMM=0
|
||||
pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'advanced_model' --run-level 'advanced_model'
|
||||
"
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
|
||||
.buildkite/common/scripts/run_cov_split.sh \
|
||||
--model-id qwen3_tts \
|
||||
--offline tests/e2e/offline_inference/test_qwen3_tts_base.py \
|
||||
--online tests/e2e/online_serving/test_qwen3_tts_base.py \
|
||||
-- -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Omni · Qwen3-Omni Test"
|
||||
@@ -124,6 +120,10 @@ steps:
|
||||
- tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- tests/e2e/online_serving/test_qwen3_omni.py
|
||||
- vllm_omni/model_executor/models/qwen3_omni/
|
||||
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
|
||||
- vllm_omni/model_executor/models/common/snake_activation.py
|
||||
- vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
|
||||
- vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
|
||||
- vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
timeout_in_minutes: 50
|
||||
@@ -135,6 +135,11 @@ steps:
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_minicpmo_4_5.py
|
||||
- tests/e2e/online_serving/test_minicpmo_4_5.py
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/models/step_audio2/
|
||||
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
|
||||
- vllm_omni/model_executor/models/cosyvoice3/utils.py
|
||||
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
|
||||
- vllm_omni/deploy/minicpmo_4_5.yaml
|
||||
timeout_in_minutes: 50
|
||||
commands:
|
||||
@@ -156,62 +161,35 @@ steps:
|
||||
- tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
|
||||
- vllm_omni/deploy/minicpmo_4_5.yaml
|
||||
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
|
||||
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
|
||||
- vllm_omni/experimental/fullduplex/
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/models/step_audio2/
|
||||
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
|
||||
- vllm_omni/model_executor/models/cosyvoice3/utils.py
|
||||
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
|
||||
- vllm_omni/core/sched/
|
||||
- vllm_omni/outputs/multimodal_accumulation.py
|
||||
- vllm_omni/engine/stage_init_utils.py
|
||||
- vllm_omni/engine/async_omni_engine.py
|
||||
- vllm_omni/engine/orchestrator.py
|
||||
- vllm_omni/entrypoints/cli/serve.py
|
||||
- vllm_omni/entrypoints/openai/api_server.py
|
||||
- vllm_omni/entrypoints/openai/duplex_capability.py
|
||||
- vllm_omni/config/stage_config.py
|
||||
timeout_in_minutes: 50
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "TTS · VoxCPM2 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_voxcpm2_tts.py
|
||||
- tests/e2e/online_serving/test_voxcpm2_tts.py
|
||||
- vllm_omni/model_executor/models/voxcpm2/
|
||||
- vllm_omni/deploy/voxcpm2.yaml
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Diffusion · Qwen Image Edit Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_edit.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
mirror_hardwares: h100_1
|
||||
- label: "Diffusion · Bagel Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_bagel.py
|
||||
- tests/e2e/online_serving/test_bagel.py
|
||||
- vllm_omni/model_executor/models/bagel/
|
||||
- vllm_omni/diffusion/models/bagel/
|
||||
- vllm_omni/model_executor/stage_input_processors/bagel.py
|
||||
- vllm_omni/deploy/bagel.yaml
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_bagel.py tests/e2e/online_serving/test_bagel.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Diffusion · Qwen Image Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
- vllm_omni/diffusion/models/dmd2/
|
||||
- vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c '
|
||||
@@ -220,23 +198,12 @@ steps:
|
||||
'
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Diffusion · Qwen Image Layered Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_layered.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_wan22_t2v.py
|
||||
- tests/e2e/online_serving/test_wan22_t2v.py
|
||||
- vllm_omni/diffusion/models/wan2_2/
|
||||
- vllm_omni/diffusion/models/dmd2/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -244,21 +211,3 @@ steps:
|
||||
pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
|
||||
"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Diffusion · Z Image Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_z_image.py
|
||||
- tests/e2e/online_serving/test_z_image.py
|
||||
- vllm_omni/diffusion/models/z_image/
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_z_image.py tests/e2e/online_serving/test_z_image.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
|
||||
mirror_hardwares: l4_1
|
||||
- label: "Diffusion · FLUX.2 Klein Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_flux2_klein.py
|
||||
- vllm_omni/diffusion/models/flux2_klein/
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
|
||||
mirror_hardwares: l4_1
|
||||
@@ -159,19 +159,7 @@ steps:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- export VLLM_USE_DEEP_GEMM="0"
|
||||
- export VLLM_MOE_USE_DEEP_GEMM="0"
|
||||
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: ":full_moon: TTS · Accuracy Test · VoxCPM2"
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
|
||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts_expansion.py -m 'full_model and cuda and tts' --run-level 'full_model'
|
||||
"
|
||||
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: ":full_moon: TTS · Perf Test"
|
||||
@@ -188,35 +176,6 @@ steps:
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: TTS · VoxCPM2 · Perf Test"
|
||||
key: nightly-tts-performance-voxcpm2
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_voxcpm2.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: TTS · Higgs-Audio-V3 · Perf Test"
|
||||
key: nightly-tts-performance-higgs-audio-v3
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export VLLM_USE_DEEP_GEMM="0"
|
||||
- export VLLM_MOE_USE_DEEP_GEMM="0"
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_higgs_audio_v3.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- group: ":card_index_dividers: Tiny Model Tests [Multi-GPU]"
|
||||
key: nightly-tiny-model-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
@@ -246,8 +205,7 @@ steps:
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2iat-test"
|
||||
steps:
|
||||
# Use explicit file shards so single-card cases do not reserve 2+ GPUs,
|
||||
# and Bagel feature coverage is not duplicated via a broad tests/e2e sweep.
|
||||
# Use explicit file shards so single-card cases do not reserve 2+ GPUs.
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Single-GPU"
|
||||
timeout_in_minutes: 120
|
||||
@@ -255,15 +213,6 @@ steps:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_qwen_image_expansion.py
|
||||
tests/e2e/online_serving/test_qwen_image_edit_expansion.py
|
||||
tests/e2e/online_serving/test_qwen_image_layered_expansion.py
|
||||
tests/e2e/online_serving/test_boogu_image.py
|
||||
tests/e2e/online_serving/test_boogu_image_edit.py
|
||||
tests/e2e/online_serving/test_boogu_image_expansion.py
|
||||
tests/e2e/online_serving/test_boogu_image_edit_expansion.py
|
||||
tests/e2e/online_serving/test_bagel_expansion.py
|
||||
tests/e2e/online_serving/test_krea2_expansion.py
|
||||
tests/e2e/offline_inference/test_krea2_expansion.py
|
||||
-m "full_model and diffusion and H100 and not distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_1
|
||||
@@ -274,47 +223,14 @@ steps:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_qwen_image_expansion.py
|
||||
tests/e2e/online_serving/test_qwen_image_edit_expansion.py
|
||||
tests/e2e/offline_inference/test_krea2_expansion.py
|
||||
-m "full_model and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Layered/LongCat/FLUX"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_qwen_image_layered_expansion.py
|
||||
-m "full_model and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Bagel"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_bagel_expansion.py
|
||||
-m "full_model and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · GLM-Image Function Test with H100"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv
|
||||
tests/e2e/online_serving/test_glm_image_expansion.py
|
||||
tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
|
||||
-m "full_model and diffusion and H100"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
- pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
mirror_hardwares: l4_4
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Doc Test"
|
||||
@@ -327,7 +243,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export VLLM_HTTP_TIMEOUT_KEEP_ALIVE=120
|
||||
- pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_qwen_image.py --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2i_matches_diffusers' -m full_model --run-level full_model
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
@@ -355,70 +271,6 @@ steps:
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Layered"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-layered
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · BAGEL"
|
||||
key: nightly-diffusion-x2iat-performance-bagel
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.5.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- group: ":card_index_dividers: Diffusion X2V Model Test"
|
||||
key: nightly-diffusion-x2v-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
@@ -440,12 +292,6 @@ steps:
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda" --run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · LingBot Video MoE Function Test"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_lingbot_video_moe.py -k "video_generation" -m "full_model and cuda" --run-level "full_model"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Doc Test"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
@@ -456,7 +302,6 @@ steps:
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx2_expansion.py -m "full_model and cuda" --run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test"
|
||||
@@ -477,8 +322,6 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_one_stage_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Perf Test"
|
||||
@@ -556,16 +399,6 @@ steps:
|
||||
- pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: ":bar_chart: Testcase Statistics"
|
||||
key: nightly-testcase-statistics
|
||||
timeout_in_minutes: 120
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
# NOTE: ":email: Nightly Collection & Email" is deprecated and will be removed soon.
|
||||
# Do not add new depends_on entries or extend this step.
|
||||
- label: ":email: Nightly Collection & Email"
|
||||
@@ -576,32 +409,20 @@ steps:
|
||||
- nightly-omni-performance-vllm-text
|
||||
- nightly-omni-performance-multi-replicas
|
||||
- nightly-tts-performance
|
||||
- nightly-tts-performance-higgs-audio-v3
|
||||
- nightly-diffusion-x2iat-performance-qwen-image
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-layered
|
||||
- nightly-diffusion-x2iat-performance-bagel
|
||||
- nightly-diffusion-x2v-performance
|
||||
- nightly-testcase-statistics
|
||||
if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
|
||||
commands:
|
||||
- pip install openpyxl
|
||||
- export DEFAULT_INPUT_DIR=tests/dfx/perf/results
|
||||
- export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-higgs-audio-v3
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
|
||||
- python tools/nightly/generate_nightly_perf_excel.py
|
||||
- python tools/nightly/generate_nightly_perf_html.py
|
||||
- python tools/nightly/send_nightly_email.py --report-file "tests/dfx/perf/results/*.xlsx, tests/dfx/perf/results/*.html"
|
||||
|
||||
@@ -169,6 +169,10 @@ steps:
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen3_omni.py
|
||||
- vllm_omni/model_executor/models/qwen3_omni/
|
||||
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
|
||||
- vllm_omni/model_executor/models/common/snake_activation.py
|
||||
- vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
|
||||
- vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
|
||||
- vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
commands:
|
||||
@@ -181,6 +185,11 @@ steps:
|
||||
- label: "Omni · MiniCPM-o 4.5 Offline Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_minicpmo_4_5.py
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/models/step_audio2/
|
||||
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
|
||||
- vllm_omni/model_executor/models/cosyvoice3/utils.py
|
||||
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
|
||||
- vllm_omni/deploy/minicpmo_4_5.yaml
|
||||
commands:
|
||||
- |
|
||||
@@ -192,6 +201,11 @@ steps:
|
||||
- label: "Omni · MiniCPM-o 4.5 Online Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_minicpmo_4_5.py
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/models/step_audio2/
|
||||
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
|
||||
- vllm_omni/model_executor/models/cosyvoice3/utils.py
|
||||
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
|
||||
- vllm_omni/deploy/minicpmo_4_5.yaml
|
||||
commands:
|
||||
- |
|
||||
@@ -210,16 +224,24 @@ steps:
|
||||
- tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
|
||||
- vllm_omni/deploy/minicpmo_4_5.yaml
|
||||
- tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
|
||||
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
|
||||
- vllm_omni/experimental/fullduplex/
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/models/step_audio2/
|
||||
- vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
|
||||
- vllm_omni/model_executor/models/cosyvoice3/utils.py
|
||||
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
|
||||
- vllm_omni/core/sched/
|
||||
- vllm_omni/outputs/multimodal_accumulation.py
|
||||
- vllm_omni/engine/stage_init_utils.py
|
||||
- vllm_omni/engine/async_omni_engine.py
|
||||
- vllm_omni/engine/orchestrator.py
|
||||
- vllm_omni/entrypoints/cli/serve.py
|
||||
- vllm_omni/entrypoints/openai/api_server.py
|
||||
- vllm_omni/entrypoints/openai/duplex_capability.py
|
||||
- vllm_omni/config/stage_config.py
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
@@ -239,25 +261,13 @@ steps:
|
||||
- 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
|
||||
- vllm_omni/model_executor/models/voxcpm2/
|
||||
- vllm_omni/deploy/voxcpm2.yaml
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "TTS · Qwen3-TTS CustomVoice Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen3_tts_customvoice.py
|
||||
- vllm_omni/model_executor/models/qwen3_tts/
|
||||
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
|
||||
- vllm_omni/model_executor/models/common/snake_activation.py
|
||||
- vllm_omni/model_executor/models/common/alias_free_activation.py
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
|
||||
- vllm_omni/deploy/qwen3_tts.yaml
|
||||
commands:
|
||||
@@ -272,6 +282,9 @@ steps:
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen3_tts_base.py
|
||||
- vllm_omni/model_executor/models/qwen3_tts/
|
||||
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
|
||||
- vllm_omni/model_executor/models/common/snake_activation.py
|
||||
- vllm_omni/model_executor/models/common/alias_free_activation.py
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
|
||||
- vllm_omni/deploy/qwen3_tts.yaml
|
||||
commands:
|
||||
@@ -284,52 +297,11 @@ steps:
|
||||
"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "TTS · Higgs-Audio-V3 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_higgs_audio_v3.py
|
||||
- vllm_omni/model_executor/models/higgs_audio_v3/
|
||||
- vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py
|
||||
- vllm_omni/deploy/higgs_multimodal_qwen3.yaml
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_USE_DEEP_GEMM=0
|
||||
export VLLM_MOE_USE_DEEP_GEMM=0
|
||||
pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Diffusion · Qwen Image Edit Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_edit.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Diffusion · Bagel Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_bagel.py
|
||||
- vllm_omni/model_executor/models/bagel/
|
||||
- vllm_omni/diffusion/models/bagel/
|
||||
- vllm_omni/model_executor/stage_input_processors/bagel.py
|
||||
- vllm_omni/deploy/bagel.yaml
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_bagel.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_wan22_t2v.py
|
||||
- vllm_omni/diffusion/models/wan2_2/
|
||||
- vllm_omni/diffusion/models/dmd2/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -342,6 +314,9 @@ steps:
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_cosmos3.py
|
||||
- vllm_omni/diffusion/models/cosmos3/
|
||||
- vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py
|
||||
- vllm_omni/experimental/world_models/
|
||||
- vllm_omni/model_extras/cosmos3.py
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -354,6 +329,8 @@ steps:
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
- vllm_omni/diffusion/models/dmd2/
|
||||
- vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -361,15 +338,3 @@ steps:
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Diffusion · Qwen Image Layered Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_layered.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
mirror_hardwares: h100_1
|
||||
@@ -30,14 +30,6 @@ steps:
|
||||
- pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: "Reliability Test - voxcpm2"
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-weekly-pipeline
|
||||
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
|
||||
commands:
|
||||
- pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- group: "Reliability Test - Invalid parameters Test"
|
||||
depends_on: upload-weekly-pipeline
|
||||
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
|
||||
@@ -108,6 +100,22 @@ steps:
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
# Own step: the image has no NanoCodec, so this file import-skips out of
|
||||
# the mark-selected TTS steps above. The overrides keep nemo_toolkit[tts]
|
||||
# from downgrading transformers; see
|
||||
# examples/offline_inference/text_to_speech/README.md.
|
||||
- label: "TTS · Gepard-1.0 · L4"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0" > /tmp/gepard-overrides.txt
|
||||
- uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" "nemo_toolkit[tts]==2.7.3" --overrides /tmp/gepard-overrides.txt
|
||||
- >-
|
||||
pytest -s -v tests/e2e/offline_inference/test_gepard_tts.py
|
||||
-m "tts and cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Diffusion · L4"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
@@ -133,4 +141,4 @@ steps:
|
||||
pytest -sv
|
||||
tests/e2e/ -m "slow and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
mirror_hardwares: h100_4
|
||||
@@ -87,8 +87,9 @@ time timeout -k 30 30m docker run \
|
||||
echo $ZE_AFFINITY_MASK
|
||||
pip install tblib==3.1.0
|
||||
cd /workspace/vllm-omni
|
||||
pytest -v -s -m "core_model and xpu and B60"
|
||||
XPU_TEST_PATHS="tests/diffusion tests/dfx tests/e2e"
|
||||
pytest -v -s $XPU_TEST_PATHS -m "core_model and xpu and B60"
|
||||
pytest -v -s tests/diffusion/quantization/test_mxfp8_config.py
|
||||
pytest -v -s -m "advanced_model and xpu and B60"
|
||||
pytest -v -s -m "omni and xpu and B60"
|
||||
pytest -v -s $XPU_TEST_PATHS -m "advanced_model and xpu and B60"
|
||||
pytest -v -s $XPU_TEST_PATHS -m "omni and xpu and B60"
|
||||
'
|
||||
@@ -69,6 +69,12 @@ steps:
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
commands:
|
||||
# Pin the Seed-TTS scorer to the second card: _get_eval_device() otherwise
|
||||
# auto-picks npu:0, which the omni server under test already fills with its
|
||||
# KV cache, so Whisper OOMs on `.to(device)` and every clip reports
|
||||
# "ASR / WER failed".
|
||||
- export SEED_TTS_WER_EVAL=1
|
||||
- export SEED_TTS_EVAL_DEVICE=npu:1
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py -m "full_model and npu and A3" --run-level full_model
|
||||
@@ -96,33 +102,6 @@ steps:
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
- group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
|
||||
key: nightly-diffusion-x2iat-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: >-
|
||||
build.env("NIGHTLY") == "1" ||
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2iat-test"
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: a3_npu_8
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- 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
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
- group: ":card_index_dividers: Diffusion I2I Model Test"
|
||||
key: nightly-HunyuanImage3-x2i-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
|
||||
@@ -1,4 +1,15 @@
|
||||
steps:
|
||||
- label: "Platforms Test"
|
||||
key: ready-platforms-test
|
||||
depends_on: upload-ready-pipeline
|
||||
mirror_hardwares: a3_npu_2
|
||||
commands:
|
||||
- >-
|
||||
pytest -s -v
|
||||
tests/platforms/npu/test_minicpmo_code2wav_npugraph.py
|
||||
-m 'core_model and npu and A3 and omni'
|
||||
--run-level 'core_model'
|
||||
|
||||
- group: ":card_index_dividers: Omni Model Test"
|
||||
key: ready-omni-group
|
||||
depends_on: upload-ready-pipeline
|
||||
|
||||
@@ -46,6 +46,8 @@ htmlcov/
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
coverage-*.xml
|
||||
coverage-*.xml.gz
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
@@ -252,6 +254,7 @@ outputs/
|
||||
results/
|
||||
generated/
|
||||
output_*/
|
||||
/tests/e2e/accuracy/artifacts/
|
||||
|
||||
# Configuration overrides
|
||||
configs/local.yaml
|
||||
|
||||
@@ -66,7 +66,7 @@ repos:
|
||||
entry: python tools/pre_commit/check_tts_adapter.py
|
||||
language: python
|
||||
pass_filenames: false
|
||||
files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$
|
||||
files: ^(vllm_omni/entrypoints/openai/serving_speech\.py|tools/pre_commit/check_tts_adapter\.py)$
|
||||
# The budget-is-now-slack reminder is printed on a PASSING run, and
|
||||
# pre-commit hides output from passing hooks unless the hook is verbose.
|
||||
verbose: true
|
||||
|
||||
@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
|
||||
---
|
||||
|
||||
*Latest News* 🔥
|
||||
- [2026/08] [VeRL-Omni](https://github.com/verl-project/verl-omni) `v0.2.0` is released: faster diffusion RL powered by vLLM-Omni (request-level/step-wise batching with FA3), rebuilt Qwen3-Omni multimodal training (DPO & GSPO), plus LTX-2.3, Qwen-Image-Edit support and more. See the [release notes](https://github.com/verl-project/verl-omni/releases/tag/v0.2.0).
|
||||
- [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support.
|
||||
- [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
|
||||
- [2026/06] Starting with [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0), vLLM-Omni publishes a stable release aligned with every even-numbered upstream vLLM minor version. [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0), [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0), [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0), and [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) continued this cadence, expanding omni and world-model support with [NVIDIA Cosmos3](recipes/cosmos3/Cosmos3-Nano.md) and DreamZero, adding models such as MiniCPM-o 4.5, MOSS-TTS, and Lance, and advancing TTS, diffusion, distributed execution, quantization, RL integration through [VeRL-Omni](https://github.com/verl-project/verl-omni), and CUDA/ROCm/MUSA/NPU/XPU coverage.
|
||||
|
||||
@@ -1,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,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,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,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
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.26.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.27.0
|
||||
FROM ${BASE_IMAGE} AS base
|
||||
|
||||
# Declare a variable to know if we want to use the nightly build or the stable build.
|
||||
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
|
||||
# we should swap over to use stable release ASAP.
|
||||
# We should avoid relying on custom commits.
|
||||
ARG USE_NIGHTLY_BUILD=0
|
||||
ARG VLLM_VERSION_OR_COMMIT_HASH=v0.26.0
|
||||
ARG VLLM_VERSION_OR_COMMIT_HASH=v0.27.0
|
||||
ARG ARG_PYTORCH_ROCM_ARCH
|
||||
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
|
||||
|
||||
@@ -79,9 +79,16 @@ RUN printf '%s\n' \
|
||||
' os._exit(_exit_code)' \
|
||||
> ${COMMON_WORKDIR}/vllm-omni/conftest.py
|
||||
|
||||
# Some optional dependencies pull in the CPU ONNX Runtime distribution. The
|
||||
# CPU and ROCm wheels own the same import path and cannot coexist, so remove
|
||||
# the CPU wheel and reinstall the ROCm wheel after dependency resolution.
|
||||
RUN cd ${COMMON_WORKDIR}/vllm-omni && \
|
||||
uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" \
|
||||
--no-cache-dir ".[dev]" --no-build-isolation && \
|
||||
python3 -m pip uninstall -y onnxruntime && \
|
||||
python3 -m pip install --no-cache-dir --force-reinstall --no-deps \
|
||||
"onnxruntime-rocm>=1.22.2" && \
|
||||
python3 -c "import importlib.metadata as m; import onnxruntime as ort; assert 'ROCMExecutionProvider' in ort.get_available_providers(), ort.get_available_providers(); assert m.version('onnxruntime-rocm'); assert not any(d.metadata['Name'] == 'onnxruntime' for d in m.distributions())" && \
|
||||
python3 -c "from torchcodec.decoders import VideoDecoder"
|
||||
|
||||
RUN ln -sf /usr/bin/python3 /usr/bin/python
|
||||
|
||||
+22
-4
@@ -43,6 +43,7 @@ nav:
|
||||
- Overview: features/README.md
|
||||
- Runtime and Stage Execution:
|
||||
- Execution Modes and Streaming: user_guide/diffusion/execution_modes.md
|
||||
- Prefill-Decode Disaggregation (experimental): features/pd_disaggregation.md
|
||||
- Sleep Mode: features/sleep_mode.md
|
||||
- Quantization:
|
||||
- Overview: user_guide/quantization/overview.md
|
||||
@@ -60,7 +61,11 @@ nav:
|
||||
- Overview and Compatibility:
|
||||
- Overview: user_guide/diffusion_features.md
|
||||
- Feature Compatibility: user_guide/feature_compatibility.md
|
||||
- CPU Offloading: user_guide/diffusion/cpu_offload.md
|
||||
- CPU Offloading:
|
||||
- Overview: user_guide/diffusion/cpu_offload.md
|
||||
- Model-Level: user_guide/diffusion/offloader/cpu_offload.md
|
||||
- Layerwise: user_guide/diffusion/offloader/layerwise_offload.md
|
||||
- Distributed Layerwise: user_guide/diffusion/offloader/distributed_layerwise_offload.md
|
||||
- Cache Acceleration:
|
||||
- TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
|
||||
- Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
|
||||
@@ -72,7 +77,13 @@ nav:
|
||||
- Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
|
||||
- Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
|
||||
- VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md
|
||||
- Attention Backends: user_guide/diffusion/attention_backends.md
|
||||
- Attention Backends:
|
||||
- Overview: user_guide/diffusion/attention_backends.md
|
||||
- Dense Backends: user_guide/diffusion/attention_backends/dense_backends.md
|
||||
- TRTLLM: user_guide/diffusion/attention_backends/trtllm.md
|
||||
- SageAttention: user_guide/diffusion/attention_backends/sage.md
|
||||
- Hugging Face Hub: user_guide/diffusion/attention_backends/huggingface_hub.md
|
||||
- RainFusion: user_guide/diffusion/attention_backends/rainfusion.md
|
||||
- Regional Compilation: user_guide/diffusion/regional_compilation.md
|
||||
- Frame Interpolation: user_guide/diffusion/frame_interpolation.md
|
||||
- Startup and Loading: user_guide/diffusion/startup_and_loading.md
|
||||
@@ -106,6 +117,7 @@ nav:
|
||||
- Async Diffusion Output: design/feature/async_diffusion_output.md
|
||||
- Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
|
||||
- Automatic Prefix Caching: design/feature/prefix_caching.md
|
||||
- Realtime AR-Diffusion: design/feature/realtime_ar_diffusion.md
|
||||
- Communication:
|
||||
- OmniConnector Implementations:
|
||||
- Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md
|
||||
@@ -125,11 +137,16 @@ nav:
|
||||
- Tensor Parallel: design/feature/tensor_parallel.md
|
||||
- VAE Patch Parallelism: design/feature/vae_parallel.md
|
||||
- Attention Optimization:
|
||||
- Backend Selection: design/feature/attention_backend_selection.md
|
||||
- Skip-Softmax: design/feature/skip_softmax.md
|
||||
- CPU Offloading:
|
||||
- Overview: design/feature/offloader/README.md
|
||||
- Model-Level: design/feature/offloader/cpu_offload.md
|
||||
- Layerwise: design/feature/offloader/layerwise_offload.md
|
||||
- Distributed Layerwise: design/feature/offloader/distributed_layerwise_offload.md
|
||||
- Cache-DiT: design/feature/cache_dit.md
|
||||
- TeaCache: design/feature/teacache.md
|
||||
- Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md
|
||||
- Distributed Layerwise Offload: design/feature/distributed_layerwise_offload.md
|
||||
- Infrastructure and Performance:
|
||||
- Prometheus Metrics: design/metrics.md
|
||||
- Speech Generation Performance Optimizations: design/qwen3_omni_tts_performance_optimization.md
|
||||
@@ -163,7 +180,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
|
||||
@@ -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.
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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 +0,0 @@
|
||||
# Meetups
|
||||
@@ -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!
|
||||
@@ -14,10 +14,6 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
|
||||
|
||||
- **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
|
||||
|
||||
## Multi-Stage Recipes
|
||||
|
||||
- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
|
||||
|
||||
## Optimization Features
|
||||
|
||||
- **[Diffusion Features Overview](../user_guide/diffusion_features.md)** - Complete overview of all diffusion model features and supported models
|
||||
@@ -1,171 +0,0 @@
|
||||
# Prefill-Decode (PD) Disaggregation
|
||||
|
||||
PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
|
||||
stages so prompt processing and token generation can run on different workers.
|
||||
|
||||
This is documented as a stage-config recipe instead of a bundled YAML because the
|
||||
deployment-specific values usually change per environment:
|
||||
|
||||
- GPU placement
|
||||
- `tensor_parallel_size`
|
||||
- connector backend and connector ports
|
||||
- connector IPs or bootstrap addresses
|
||||
|
||||
Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
|
||||
and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
|
||||
changes below.
|
||||
|
||||
## Requirements
|
||||
|
||||
- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
|
||||
- A KV connector supported by vLLM, such as `MooncakeConnector`
|
||||
- Matching `tensor_parallel_size` on the prefill and decode thinker stages
|
||||
|
||||
## 1. Split the thinker into prefill and decode stages
|
||||
|
||||
Replace the original thinker stage with two stages:
|
||||
|
||||
```yaml
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
is_prefill_only: true
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 16
|
||||
model_stage: thinker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.9
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 1
|
||||
kv_transfer_config:
|
||||
kv_connector: "MooncakeConnector"
|
||||
kv_role: "kv_producer"
|
||||
kv_rank: 0
|
||||
kv_parallel_size: 2
|
||||
kv_connector_extra_config:
|
||||
mooncake_bootstrap_port: 25201
|
||||
final_output: false
|
||||
is_comprehension: true
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
top_k: 1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
is_decode_only: true
|
||||
runtime:
|
||||
devices: "1"
|
||||
engine_args:
|
||||
max_num_seqs: 64
|
||||
model_stage: thinker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.9
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 1
|
||||
kv_transfer_config:
|
||||
kv_connector: "MooncakeConnector"
|
||||
kv_role: "kv_consumer"
|
||||
kv_rank: 1
|
||||
kv_parallel_size: 2
|
||||
kv_connector_extra_config:
|
||||
mooncake_bootstrap_port: 25202
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
is_comprehension: true
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
top_k: 1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
```
|
||||
|
||||
Notes:
|
||||
|
||||
- `is_prefill_only: true` marks the thinker stage that only saves KV.
|
||||
- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
|
||||
- `kv_transfer_config` is required on both stages.
|
||||
- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
|
||||
prefill side only processes the prompt and exports KV.
|
||||
|
||||
## 2. Shift the downstream stages by one index
|
||||
|
||||
After inserting the extra thinker stage, renumber the remaining stages:
|
||||
|
||||
```yaml
|
||||
- stage_id: 2
|
||||
runtime:
|
||||
devices: "2"
|
||||
engine_input_source: [1]
|
||||
sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only
|
||||
|
||||
- stage_id: 3
|
||||
runtime:
|
||||
devices: "2"
|
||||
engine_args:
|
||||
max_num_seqs: 1
|
||||
engine_input_source: [2]
|
||||
sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_token_only
|
||||
```
|
||||
|
||||
Compared with the default Qwen3-Omni config:
|
||||
|
||||
- the talker becomes stage `2` instead of stage `1`
|
||||
- the code2wav stage becomes stage `3` instead of stage `2`
|
||||
- the talker now reads from decode stage `1`
|
||||
|
||||
## 3. Add runtime edges for the four-stage pipeline
|
||||
|
||||
```yaml
|
||||
runtime:
|
||||
enabled: true
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
- from: 1
|
||||
to: 2
|
||||
- from: 2
|
||||
to: 3
|
||||
```
|
||||
|
||||
## 4. Launch with your custom config
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--deploy-config /path/to/qwen3_omni_pd.yaml
|
||||
```
|
||||
|
||||
## Operational Notes
|
||||
|
||||
- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
|
||||
pair. Keep prefill and decode at the same `tensor_parallel_size`.
|
||||
- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
|
||||
separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
|
||||
- Choose connector ports and addresses that match your deployment. The values
|
||||
shown above are examples only.
|
||||
@@ -38,11 +38,12 @@ Each entry under `stages:` accepts any `StageDeployConfig` field directly (no ne
|
||||
|-------|------|----------|---------|-------------|
|
||||
| `stage_id` | int | required | — | Stage identity; matched against `PipelineConfig.stages[*].stage_id`. |
|
||||
| `max_num_seqs` | int | optional | `64` | Max concurrent sequences per stage. |
|
||||
| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage memory budget. |
|
||||
| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage total memory target; used for automatic KV-cache sizing. |
|
||||
| `kv_cache_memory_bytes` | int \| null | optional | `null` | Explicit per-rank KV-cache budget in bytes; overrides automatic sizing. |
|
||||
| `tensor_parallel_size` | int | optional | `1` | TP degree for this stage. |
|
||||
| `enforce_eager` | bool | optional | `false` | Disable CUDA graphs. |
|
||||
| `max_num_batched_tokens` | int | optional | `32768` | Prefill budget. |
|
||||
| `max_model_len` | int \| null | optional | `null` | Per-stage context length (auto-sets `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1` when larger than HF default). |
|
||||
| `max_num_batched_tokens` | int | optional | `32768` | Per-stage prefill/token budget; also contributes to the native maximum in-flight token limit. |
|
||||
| `max_model_len` | int \| null | optional | `null` | Per-sequence context or KV length; `-1` enables native cache-capacity auto-fitting, while values above the HF default auto-set `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1`. |
|
||||
| `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. |
|
||||
| `devices` | str | optional | `"0"` | `CUDA_VISIBLE_DEVICES`-style device list. |
|
||||
| `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_<n>`; values are names registered under top-level `connectors:`. |
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -327,6 +327,99 @@ On PR builds, `upload_pipeline.py` logs `skip '…' (no changes under …)` for
|
||||
- Use **groups** for dashboard readability; keep **E2E Test** as the group name CUDA diff filtering expects for `--e2e` nightly runs on main.
|
||||
- Prefix labels by model domain: **Omni ·**, **TTS ·**, **Diffusion ·**, **Simple ·**, etc., matching existing steps.
|
||||
|
||||
### Per-model coverage
|
||||
|
||||
Pilot (v1): one Merge-tier (L3) job in `.buildkite/cuda/test-merge.yml` uploads
|
||||
per-model, per-entry-mode coverage as Buildkite artifacts — "TTS · Qwen3-TTS Base
|
||||
Test". It runs on a single GPU so the pilot is cheap to reproduce. (A second
|
||||
diffusion pilot, "Diffusion · Bagel Test", was retired from merge; reintroduce
|
||||
another diffusion coverage job if a second pilot is needed again.) This is a
|
||||
pilot, not a rollout: no dashboard/visualization lives in this repo, and nothing
|
||||
is gated on coverage.
|
||||
|
||||
**Naming convention**: `coverage-<model_id>-<mode>-<step_id>.xml.gz`, where
|
||||
`<model_id>` is the model's directory name under
|
||||
`vllm_omni/diffusion/models/<model_id>/` or
|
||||
`vllm_omni/model_executor/models/<model_id>/` (e.g. `z_image`, `qwen3_omni`),
|
||||
`<mode>` is `online` or `offline`, and `<step_id>` is `BUILDKITE_STEP_ID`
|
||||
(`local` outside Buildkite). `gunzip` before feeding a report to a Cobertura
|
||||
consumer.
|
||||
|
||||
Reports are gzipped because `--cov=vllm_omni` sets coverage's `source`, which
|
||||
makes it walk the package and emit one `<line>` element per statement for every
|
||||
file — including the ones a single model's tests never import. That inventory,
|
||||
not the coverage, is the size: a Bagel report is 7.1 MB, of which 96.5% is
|
||||
`<line>` elements and 77% belongs to the 840 of 1169 files that run never
|
||||
imported. A run covering zero lines produces the same 7 MB. gzip takes it to
|
||||
roughly 410 KB and changes nothing about the data.
|
||||
|
||||
`<step_id>` is what keeps two steps that cover the same model apart. Artifacts are
|
||||
scoped to the build, and a nightly build carries the ready, merge and nightly tiers
|
||||
at once, so several same-model steps land in one namespace — `qwen3_tts` already has
|
||||
both a Base and a CustomVoice job, and `minicpmo_4_5` a base and a duplex one. Two
|
||||
uploads on one path do not fail loudly: an exact-name download can report the path
|
||||
as ambiguous, and a glob download fetches every match concurrently and renames each
|
||||
onto the same destination, so whichever finishes last wins. The step id remains the
|
||||
logical step identifier across retries, while Buildkite's default artifact lookup
|
||||
selects only the latest attempt — `BUILDKITE_JOB_ID` would instead give every
|
||||
attempt its own filename, leaving a failed attempt's partial report to be picked up
|
||||
alongside the real one. A step using `parallelism`/matrix would need
|
||||
`BUILDKITE_PARALLEL_JOB` added, since its jobs share a step id; neither pilot is
|
||||
one.
|
||||
|
||||
**Opting in a new model**: replace the job's combined `pytest` command with
|
||||
[`run_cov_split.sh`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/scripts/run_cov_split.sh),
|
||||
which runs each mode, names the reports, uploads them, and fails the step if
|
||||
either half or the upload failed:
|
||||
|
||||
```yaml
|
||||
commands:
|
||||
- |
|
||||
.buildkite/common/scripts/run_cov_split.sh \
|
||||
--model-id <model_id> \
|
||||
--offline <offline test path> \
|
||||
--online <online test path> \
|
||||
-- <the job's existing pytest flags, e.g. -m '...' --run-level '...'>
|
||||
```
|
||||
|
||||
Everything after `--` is forwarded verbatim to both runs with quoting preserved,
|
||||
so an expression like `-m 'advanced_model and cuda'` survives intact and a job can
|
||||
pass anything else it needs. `--offline`/`--online` are repeatable and only one is
|
||||
required, so a job with tests in a single mode runs just that half. Runtime is
|
||||
bounded solely by the step's `timeout_in_minutes` — raise it, since the split
|
||||
loads the model once per mode.
|
||||
|
||||
Before splitting, confirm each half actually collects at least one test under the job's forwarded filters — pytest exits with code 5 ("no tests collected") if a half is empty, which the script reports as a failure, red-ing a job that used to pass as one combined invocation.
|
||||
|
||||
Also check the tests' `num_cards` against the job's `mirror_hardwares` preset.
|
||||
`cuda_marks` turns `num_cards > 1` into `skipif(device_count() < num_cards)`, so a
|
||||
test asking for more GPUs than the preset provides is silently skipped and its
|
||||
coverage artifact comes back empty. Splitting such a job produces a file that
|
||||
looks fine and measures nothing — "Diffusion · Z Image Test" is the current
|
||||
example (online tests want 4 cards, the job runs on `l4_1`), which is why it is
|
||||
not a coverage pilot.
|
||||
|
||||
Online-mode coverage depends on the `[tool.coverage.run]` `patch`/`sigterm`
|
||||
settings in `pyproject.toml`: online tests launch the server with
|
||||
`subprocess.Popen` and stop it with SIGTERM, and without those settings the XML
|
||||
reflects only the pytest parent process, not the server's code paths.
|
||||
|
||||
The upload depends on `buildkite-agent` being callable inside the container. The
|
||||
`kubernetes` presets (`h100_*`, `*_npu_*`) provide it; the `docker` ones (`l4_*`)
|
||||
only do because they set `mount-buildkite-agent: true`. A new docker preset that
|
||||
runs a coverage job needs the same.
|
||||
|
||||
List both `run_cov_split.sh` and `pyproject.toml` in every opted-in job's
|
||||
`source_file_dependencies` — both change what the job measures, so without them a
|
||||
change there is filtered out of normal PR builds and only surfaces in a later
|
||||
nightly. `tests/buildkite/test_upload_pipeline.py` covers the filter behavior with
|
||||
a synthetic job (it does not pin real merge labels).
|
||||
Editing only the surrounding CI YAML still does not schedule them, so a PR that
|
||||
touches just the wiring needs a full E2E run (or the commands run on a GPU host)
|
||||
to produce artifacts. When checking a new model's
|
||||
artifacts, compare `lines-covered` between the online and offline XML rather than
|
||||
just confirming both files exist.
|
||||
|
||||
### Validation checklist
|
||||
|
||||
| Check | Command / location |
|
||||
|
||||
@@ -184,9 +184,6 @@ Failed jobs: 1/2
|
||||
# Reliability: HunyuanImage DiT (H100 × 4)
|
||||
pytest -s -v dfx/reliability/test_reliability_hunyuan_image.py -m slow
|
||||
|
||||
# Reliability: VoxCPM2 (L4)
|
||||
pytest -s -v dfx/reliability/test_reliability_voxcpm2.py -m slow
|
||||
|
||||
```
|
||||
|
||||
The latest L5 CI jobs (reliability + invalid-parameter weekly steps) are in [test-weekly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-weekly.yml).
|
||||
|
||||
@@ -562,7 +562,6 @@ L5 level testing focuses on the performance of model services under ***long-runn
|
||||
- `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default).
|
||||
- `test_reliability_wan22.py` — Wan2.2 I2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
|
||||
- `test_reliability_hunyuan_image.py` — HunyuanImage-3.0-Instruct DiT-only reliability (`/v1/images/generations` under OOM and process kill; deploy `hunyuan_image3_dit.yaml`, H100 × 4).
|
||||
- `test_reliability_voxcpm2.py` — VoxCPM2 TTS reliability (`/v1/audio/speech` under OOM and process kill; deploy `voxcpm2.yaml`, L4 × 1).
|
||||
- `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`).
|
||||
- `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready.
|
||||
- `README.md` — Short local run commands for this directory.
|
||||
@@ -662,13 +661,12 @@ Reliability tests are **short fault-injection** integration runs (L5 **(b)** in
|
||||
| `test_reliability_qwen3_omni.py` | Qwen3-Omni: OOM vs **text vs text+audio** error contract, large multimodal chat under OOM, concurrent pressure, **SIGKILL** on `VLLM::Worker`, `/health` → 503 + fast-fail + concurrent chat; optional OOM recovery scenario (may be skipped while tracked in issues). |
|
||||
| `test_reliability_wan22.py` | Wan2.2 I2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
|
||||
| `test_reliability_hunyuan_image.py` | HunyuanImage DiT-only: large `/v1/images/generations` under OOM, **SIGKILL** on `vLLM-Omni::` workers and serve/tree targets, health / fast-fail / concurrent image requests; some OOM/recovery cases may be skipped while tracked in issues. |
|
||||
| `test_reliability_voxcpm2.py` | VoxCPM2: `/v1/audio/speech` under OOM (error contract), **SIGKILL** on `VLLM::` workers and serve/tree targets, health / fast-fail / concurrent speech requests; some OOM cases may be skipped while tracked in issues. |
|
||||
| `README.md` | Minimal run / collect examples. |
|
||||
|
||||
##### Parametrization and markers
|
||||
|
||||
- Each test module defines a **`RELIABILITY_SCENARIOS`** list (`test_name`, `server_params`: model, `stage_config_name` or diffusion `server_args`, etc.). **`create_reliability_omni_server_params()`** in `tests/dfx/conftest.py` resolves stage paths (including XPU substitutions where applicable) and builds **`OmniServerParams`** lists consumed by **`@pytest.mark.parametrize(..., indirect=True)`** on `omni_server` or `omni_server_function`.
|
||||
- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** or **`@hardware_test(res={"cuda": "L4"}, num_cards=1)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100; VoxCPM2 **1**× L4).
|
||||
- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100).
|
||||
- **Process-kill** tests use **`@pytest.mark.skipif(os.name == "nt", ...)`** because injection uses POSIX **`pgrep` / `kill`**.
|
||||
|
||||
##### CI trigger
|
||||
@@ -680,13 +678,11 @@ Weekly Buildkite (`.buildkite/cuda/test-weekly.yml`) runs one step per model sui
|
||||
| Reliability Test - qwen3-omni | `test_reliability_qwen3_omni.py` | H100 × 2 (`mithril-h100-pool`) |
|
||||
| Reliability Test - wan22 | `test_reliability_wan22.py` | H100 × 2 (`mithril-h100-pool`) |
|
||||
| Reliability Test - hunyuan-image | `test_reliability_hunyuan_image.py` | H100 × 4 (`mithril-h100-pool`) |
|
||||
| Reliability Test - voxcpm2 | `test_reliability_voxcpm2.py` | L4 × 1 (`gpu_1_queue`) |
|
||||
|
||||
```bash
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
|
||||
```
|
||||
|
||||
##### Local commands
|
||||
@@ -696,7 +692,6 @@ pytest --collect-only tests/dfx/reliability
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
|
||||
```
|
||||
|
||||
##### Adding a new model suite
|
||||
@@ -709,4 +704,4 @@ pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
|
||||
</details>
|
||||
|
||||
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_<model>.py` alongside a matching JSON config)
|
||||
- ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_<model>.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`, `voxcpm2`; add `test_reliability_<suite>.py` and a matching step in `.buildkite/cuda/test-weekly.yml` for new models)
|
||||
- ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_<model>.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`; add `test_reliability_<suite>.py` and a matching step in `.buildkite/cuda/test-weekly.yml` for new models)
|
||||
@@ -4,6 +4,9 @@
|
||||
> `experimental/ar_diffusion` contract. It follows the ownership and lifecycle
|
||||
> direction in [#5137](https://github.com/vllm-project/vllm-omni/issues/5137),
|
||||
> but it is not yet part of the stable module-design hierarchy.
|
||||
>
|
||||
> The end-to-end request, identity, failure, and memory contracts are documented
|
||||
> in [Realtime AR-Diffusion sessions](feature/realtime_ar_diffusion.md).
|
||||
|
||||
`ARDiffusionEngine` is an opt-in diffusion backend for pipelines that keep
|
||||
autoregressive attention KV across requests. The engine only selects
|
||||
@@ -51,14 +54,13 @@ time, and its inherited batch/step entry points also fail explicitly. Supporting
|
||||
those paths requires a batch-aware state-binding contract rather than silently
|
||||
bypassing `bind_ar_diffusion_state()`.
|
||||
|
||||
The current implementation retains one session per model instance. A request
|
||||
for a different `session_id` closes the resident session through the same LRU
|
||||
lifecycle used by explicit close and failure cleanup before allocating the new
|
||||
one. Self-attention and cross-attention memory are therefore sized for one
|
||||
resident session even if the pipeline declares a larger `session_capacity`.
|
||||
This prioritizes rollout throughput and prevents block-pool exhaustion before a
|
||||
count-based eviction can run. Multi-session residency can be added later with a
|
||||
memory-backed capacity calculation or allocation-pressure eviction.
|
||||
The implementation may retain multiple sessions per model instance. It selects
|
||||
the largest feasible capacity up to the pipeline's `session_capacity`, based on
|
||||
self-attention pages, cross-attention allocations, scratch, and declared
|
||||
model-owned state. One session is admitted whenever it fits actual free device
|
||||
memory; `gpu_memory_fraction` controls additional residency. Requests beyond
|
||||
that capacity evict the least-recently-used session through the same lifecycle
|
||||
hook used by explicit close and failure cleanup.
|
||||
|
||||
Scratch storage is also capability-driven. `frames_per_block` reserves space for
|
||||
an uncommitted current video block, while `max_scratch_tokens_per_branch`
|
||||
|
||||
@@ -406,7 +406,7 @@ the source for configuration, compatibility, and implementation details.
|
||||
stage contract.
|
||||
* **Quantization and memory efficiency:** [Quantization](feature/quantization.md)
|
||||
resolves per-pipeline or per-component quantization configurations, while
|
||||
[Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
|
||||
[Distributed Layerwise Offload](feature/offloader/distributed_layerwise_offload.md)
|
||||
streams diffusion blocks from host memory within the existing parallel
|
||||
topology.
|
||||
|
||||
|
||||
@@ -117,7 +117,22 @@ After (async D2H):
|
||||
|
||||
### Batch Split
|
||||
|
||||
When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently.
|
||||
When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently. If `OUTPUT_READY` arrives before `execute_batch` registers the split map, `execute_batch` directly adopts the early output.
|
||||
|
||||
### Reliability, Lifecycle & Timeout Behavior
|
||||
|
||||
1. **Drain Before Memory Release (`drain_async_outputs` / `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S`)**:
|
||||
When workers transition into sleep states (`handle_sleep_task`) or invoke memory-releasing methods, `WorkerProc.drain_async_outputs()` blocks up to `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S` (10.0s) until all in-flight background D2H/SHM packing tasks finish. This acts as a synchronization barrier preventing worker sleep cycles from releasing device tensors while the side CUDA stream is still actively reading them. Draining typically completes within 10–50 ms under normal operation.
|
||||
|
||||
2. **Timeout Diagnostics (`_ASYNC_OUTPUT_TIMEOUT`)**:
|
||||
The engine waits for `wait_output_ready()` with a timeout of `_ASYNC_OUTPUT_TIMEOUT` (30.0s). If a timeout occurs, the executor logs detailed diagnostic bookkeeping (pending futures, cached outputs, and batch split maps) to identify whether the stall occurred in the worker, pump, or waiter.
|
||||
|
||||
3. **Queue Serialization & Atomic Future Resolution**:
|
||||
- `result_mq` writes are serialized with `_result_mq_lock` to protect the single-writer `MessageQueue` against concurrent writes from the worker main loop (`COMPUTE_DONE`) and background thread (`OUTPUT_READY`).
|
||||
- `ResultPumpThread` unpacks shared-memory payloads first, and then atomically resolves waiting futures or populates `_completed_outputs` under `_futures_lock`, avoiding race conditions with concurrent `wait_output_ready()` invocations.
|
||||
|
||||
4. **Storage-Aware IPC Packing**:
|
||||
`pack_diffusion_output_shm()` evaluates `max(view_bytes, storage_bytes)` against `_SHM_TENSOR_THRESHOLD` (1 MB) to ensure that per-request slices sharing a large batch storage are moved to POSIX shared memory rather than serialized through pickle over the MessageQueue inline buffer.
|
||||
|
||||
## Model Coverage
|
||||
|
||||
@@ -131,6 +146,7 @@ All models running in request-mode (`step_execution=False`, the default) automat
|
||||
|---|---|---|
|
||||
| **HunyuanImage-3.0** | Image | `False` |
|
||||
| **Qwen-Image** | Image | `True` |
|
||||
| **LTX-2.3** | Video / Audio | `False` |
|
||||
|
||||
Other models with `step_execution=False` are also supported but not yet verified.
|
||||
|
||||
|
||||
@@ -0,0 +1,89 @@
|
||||
# Diffusion Attention Backend Selection
|
||||
|
||||
This document defines the selection and extension contract for diffusion
|
||||
attention backends. User-facing backend choices, installation, and tuning are
|
||||
in the [attention backend guides](../../user_guide/diffusion/attention_backends.md).
|
||||
|
||||
## Scope
|
||||
|
||||
The contract applies to DiT and other diffusion attention layers. It is
|
||||
separate from vLLM's autoregressive attention selector.
|
||||
|
||||
The selection path has four responsibilities:
|
||||
|
||||
1. normalize user configuration into `AttentionConfig` and `AttentionSpec`;
|
||||
2. resolve a spec for an attention role;
|
||||
3. ask the active platform to validate an explicit backend or choose a
|
||||
hardware default; and
|
||||
4. load the selected `AttentionBackend` class from the registry path.
|
||||
|
||||
## Resolution contract
|
||||
|
||||
`get_attn_backend_for_role()` resolves configuration in this order:
|
||||
|
||||
1. exact `per_role[role]`;
|
||||
2. category `per_role[role_category]`;
|
||||
3. `default`;
|
||||
4. platform default.
|
||||
|
||||
An explicit resolution returns both the backend class and its
|
||||
`AttentionSpec`. A platform-default resolution returns the class and `None`.
|
||||
Layers must therefore treat the spec as optional and must not infer that a
|
||||
missing spec means SDPA.
|
||||
|
||||
Class resolution is cached by backend name, head size, and whether TRTLLM may
|
||||
be selected automatically. Logging is cached separately by role and source so
|
||||
the startup record identifies why each role received its backend.
|
||||
|
||||
## Registry and platform boundary
|
||||
|
||||
`DiffusionAttentionBackendEnum` maps stable configuration names to default
|
||||
qualified class paths. `register_diffusion_backend()` may replace a path at
|
||||
runtime without changing the public enum value.
|
||||
|
||||
The active `OmniPlatform` owns compatibility policy through
|
||||
`get_diffusion_attn_backend_cls()`. It must:
|
||||
|
||||
- validate explicit selections and fail with an actionable error when the
|
||||
requested kernel cannot run;
|
||||
- choose only an available, compatible default when no backend is explicit;
|
||||
- account for head size and other platform-visible constraints; and
|
||||
- return a qualified class path implementing `AttentionBackend`.
|
||||
|
||||
The selector must not duplicate device capability or package-availability
|
||||
policy that belongs to the platform.
|
||||
|
||||
## Typed backend options
|
||||
|
||||
Backend-specific settings remain typed fields on `AttentionSpec` rather than
|
||||
unstructured keyword dictionaries:
|
||||
|
||||
- `quant` is consumed by FlashInfer and TRTLLM with backend-specific value
|
||||
validation;
|
||||
- `skip_softmax` is consumed by TRTLLM; and
|
||||
- `block_sparse` is consumed by block-sparse backends such as RainFusion.
|
||||
|
||||
A backend reads only the fields it owns and rejects incompatible values. New
|
||||
options should be added to a shared typed spec only when more than one backend
|
||||
shares their semantics; otherwise add a dedicated typed block.
|
||||
|
||||
## Model integration contract
|
||||
|
||||
Each diffusion attention call declares a stable role and, when appropriate, a
|
||||
broader category. Model-specific roles permit precise overrides; categories
|
||||
preserve common `self` or `cross` policies. A model also decides whether its
|
||||
path is compatible with automatic TRTLLM selection, because only the model
|
||||
knows whether masking and packing satisfy the kernel contract.
|
||||
|
||||
## Adding or changing a backend
|
||||
|
||||
An implementation change is complete only when it:
|
||||
|
||||
1. implements the `AttentionBackend` interface and exposes a stable enum name;
|
||||
2. adds platform validation and default routing where appropriate;
|
||||
3. defines and validates any typed configuration;
|
||||
4. tests explicit selection, default selection, and incompatible paths;
|
||||
5. documents installation, fallback behavior, and quality implications in the
|
||||
corresponding user guide; and
|
||||
6. updates the overview matrix without moving algorithm details back into the
|
||||
landing page.
|
||||
@@ -1,156 +0,0 @@
|
||||
# Distributed Layerwise Offload
|
||||
|
||||
This document describes distributed layerwise offload (DLO) for diffusion
|
||||
models. DLO keeps only a small number of DiT blocks on the accelerator and
|
||||
streams the remaining blocks from host memory. The distributed backend can
|
||||
either shard those host-side weights across an existing parallel group or keep
|
||||
the standard loader's rank-local weights and avoid an additional collective.
|
||||
|
||||
For user-facing commands, see the
|
||||
[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md)
|
||||
and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
|
||||
|
||||
## Status
|
||||
|
||||
DLO is implemented for multi-device diffusion execution. The default
|
||||
AllGather path is the primary path for DP and SP deployments. The
|
||||
`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
|
||||
for standard-loader sharding, workstation bring-up, and systems where an
|
||||
additional DLO collective is undesirable, but it does not reduce host weight
|
||||
storage across ranks.
|
||||
|
||||
The compatibility matrix below describes the current implementation. The
|
||||
unit-level guards are covered, but not every parallelism combination has a
|
||||
full model-and-hardware end-to-end test.
|
||||
|
||||
## Design
|
||||
|
||||
### DLO consumes the existing parallel topology
|
||||
|
||||
DLO does not create a new DP, TP, or SP topology. It reads the configured
|
||||
`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the
|
||||
standard distributed groups have been initialized.
|
||||
|
||||
The DLO weight-sharding group is selected as follows:
|
||||
|
||||
1. Use the existing DP group when `data_parallel_size > 1`.
|
||||
2. When DP is one and SP is greater than one, use the SP group.
|
||||
3. Otherwise, run rank-locally without a DLO process group.
|
||||
|
||||
TP is deliberately not used as DLO's AllGather group. HSDP has its own
|
||||
parameter-sharding lifecycle and is not allowed to be sharded a second time by
|
||||
DLO's AllGather path.
|
||||
|
||||
### AllGather path
|
||||
|
||||
With the default `dlo_use_allgather=True`, each rank stores approximately
|
||||
`1 / group_size` of each streamable block in pinned host memory. The next
|
||||
block's shard is copied to a device buffer and reconstructed with
|
||||
`all_gather_into_tensor` on a communication stream while the current block is
|
||||
executing.
|
||||
|
||||
```text
|
||||
Compute: [Block N] [Block N+1] [Block N+2]
|
||||
H2D: [shard N+1] [shard N+2]
|
||||
AllGather: [full N+1] [full N+2]
|
||||
Buffers: [current slot] [prefetch slot] [current slot]
|
||||
```
|
||||
|
||||

|
||||
|
||||
The backend uses two shared device buffers, so accelerator weight residency is
|
||||
bounded by the largest streamed blocks rather than the complete model.
|
||||
|
||||
When DP is greater than one, the engine can process one request per DP rank in
|
||||
the same denoising wave. Because AllGather is a collective, all participating
|
||||
requests must take the same execution path at every denoising step.
|
||||
|
||||
### Rank-local path without DLO AllGather
|
||||
|
||||
With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
|
||||
one. The regular model loader remains responsible for preparing each rank's
|
||||
weights, including TP-local tensors or HSDP-managed parameters. DLO then
|
||||
streams those rank-local tensors block by block using H2D copies only.
|
||||
|
||||
This mode means:
|
||||
|
||||
- DP still provides independent replicas, but DLO does not shard weights
|
||||
across DP ranks.
|
||||
- SP still performs its normal activation/attention collectives, but DLO does
|
||||
not shard weights across SP ranks.
|
||||
- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
|
||||
DLO's additional weight AllGather is disabled.
|
||||
- Pure DP deployments keep a full host-side model copy per rank, subject to
|
||||
shared OS page-cache behavior.
|
||||
- The scheduler does not require a synchronized DP request wave for DLO.
|
||||
|
||||
This path is intentionally not implemented by reusing the AllGather mmap
|
||||
loader. It relies on the standard loader so model-specific TP/HSDP transforms
|
||||
remain intact.
|
||||
|
||||
## Parallelism compatibility
|
||||
|
||||
| Parallelism | DLO + AllGather | DLO without AllGather |
|
||||
|---|---|---|
|
||||
| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. |
|
||||
| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
|
||||
| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
|
||||
| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
|
||||
|
||||
### Combined dimensions
|
||||
|
||||
- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater
|
||||
than one; SP continues to use its own sequence-parallel group. If DP is one,
|
||||
the SP group becomes DLO's sharding group in AllGather mode.
|
||||
- **DP + TP/SP without AllGather:** standard model loading defines the
|
||||
rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight
|
||||
collective.
|
||||
- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
|
||||
DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
|
||||
materialization and synchronization.
|
||||
- **HSDP + DP or TP:** rejected independently by the diffusion parallel
|
||||
configuration.
|
||||
|
||||
## Request and loading constraints
|
||||
|
||||
AllGather DP multi-concurrency requires:
|
||||
|
||||
- explicit `num_inference_steps`;
|
||||
- the same `num_inference_steps` for all requests in a wave; and
|
||||
- identical request arguments that affect the collective execution path.
|
||||
|
||||
The no-AllGather path does not impose these DLO-specific synchronized-wave
|
||||
requirements.
|
||||
|
||||
The mmap loader is used only by the supported DLO+AllGather path when the
|
||||
model has a compatible checkpoint layout. Online quantization is incompatible
|
||||
with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
|
||||
quantization. Models that do not meet the mmap requirements use the regular
|
||||
loader path.
|
||||
|
||||
## Validation coverage
|
||||
|
||||
Current source-level validation includes:
|
||||
|
||||
- HSDP + DLO + AllGather rejection;
|
||||
- HSDP + DLO without AllGather acceptance at configuration level;
|
||||
- TP rejection in the DLO+AllGather mmap path;
|
||||
- resident-layer requests requiring no-AllGather;
|
||||
- DP request-wave validation for denoising-step compatibility;
|
||||
- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
|
||||
tests.
|
||||
|
||||
The highest-value missing coverage is end-to-end numerical comparison against
|
||||
ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather
|
||||
on the target CUDA/NCCL or CANN/HCCL hardware.
|
||||
|
||||
## Recommendations
|
||||
|
||||
- Use **DP + DLO AllGather** for the supported throughput and host-memory
|
||||
scaling path.
|
||||
- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
|
||||
not the goal.
|
||||
- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
|
||||
the expectation of higher host-memory use and lower validation confidence.
|
||||
- Prefer **HSDP alone** for production HSDP deployments until the combined
|
||||
HSDP + DLO no-AllGather path has broader end-to-end coverage.
|
||||
@@ -0,0 +1,80 @@
|
||||
# CPU Offloading
|
||||
|
||||
This document defines the shared architecture for diffusion CPU offloading.
|
||||
The three strategies have separate user and design pages:
|
||||
|
||||
| Strategy | User guide | Design contract |
|
||||
| --- | --- | --- |
|
||||
| Model-level | [Guide](../../../user_guide/diffusion/offloader/cpu_offload.md) | [Design](cpu_offload.md) |
|
||||
| Layerwise | [Guide](../../../user_guide/diffusion/offloader/layerwise_offload.md) | [Design](layerwise_offload.md) |
|
||||
| Distributed layerwise | [Guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md) | [Design](distributed_layerwise_offload.md) |
|
||||
|
||||
## Strategy selection
|
||||
|
||||
`OffloadConfig.from_od_config()` converts diffusion configuration into one
|
||||
`OffloadStrategy`. If callers set multiple flags, selection is deterministic:
|
||||
|
||||
1. `DISTRIBUTED_LAYERWISE`;
|
||||
2. `LAYERWISE`;
|
||||
3. `MODEL_LEVEL`; and
|
||||
4. `NONE`.
|
||||
|
||||
The factory logs when a higher-priority strategy disables another. It derives
|
||||
parallel and HSDP state from `DiffusionParallelConfig`; callers do not provide
|
||||
a separate offload group size.
|
||||
|
||||
`get_offload_backend()` then validates platform offload support, resolves the
|
||||
device, and creates exactly one backend. Returning `None` means offloading is
|
||||
disabled or unsupported and must not leave partially installed hooks.
|
||||
|
||||
## Shared lifecycle
|
||||
|
||||
Every backend implements `OffloadBackend`:
|
||||
|
||||
- `enable(pipeline)` discovers modules, establishes initial residency, and
|
||||
installs hooks;
|
||||
- `disable()` removes owned hooks and resources; and
|
||||
- `is_enabled()` reports lifecycle state.
|
||||
|
||||
`disable()` does not promise to restore the pipeline's original device
|
||||
placement. The caller owns any subsequent rematerialization.
|
||||
|
||||
Hooks are registered through `HookRegistry` and `ModelHook`; offload backends
|
||||
must use distinct hook names and remove only hooks they own.
|
||||
|
||||
## Discovery and topology
|
||||
|
||||
Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
|
||||
|
||||
- `_dit_modules`;
|
||||
- `_encoder_modules`;
|
||||
- `_vae_modules`; and
|
||||
- `_resident_modules`.
|
||||
|
||||
Dotted paths are supported. Legacy pipelines may use the fallback scan of
|
||||
well-known attribute names, but new integrations should declare components
|
||||
explicitly.
|
||||
|
||||
Ordinary layerwise offload obtains topology from declared
|
||||
`_layerwise_offload_blocks_attrs` (with a deprecated singular-name
|
||||
compatibility path). Distributed layerwise offload may instead use a pipeline
|
||||
`OffloadPlan`, then falls back to declared block attributes and its documented
|
||||
compatibility discovery. Discovery metadata describes structure only; the
|
||||
backend remains responsible for transfer, synchronization, and storage
|
||||
ownership.
|
||||
|
||||
## Cross-strategy invariants
|
||||
|
||||
- At most one strategy owns offload hooks for a pipeline.
|
||||
- A parameter has one authoritative host representation while offloaded.
|
||||
- Device storage is not freed until dependent compute or transfer work is
|
||||
complete.
|
||||
- Non-persistent and model-specific buffers remain correct after movement or
|
||||
checkpoint rematerialization.
|
||||
- Unsupported parallel or loading combinations fail before hooks mutate the
|
||||
model.
|
||||
- Platform streams, events, synchronization, and cache management go through
|
||||
the vLLM-Omni platform abstraction.
|
||||
|
||||
The diffusion [Offloader module design](../../module/diffusion/offloader.md)
|
||||
describes how these feature contracts fit into the larger diffusion runtime.
|
||||
@@ -0,0 +1,56 @@
|
||||
# Model-Level Offload
|
||||
|
||||
Model-level offload enforces mutual exclusion between pipeline component
|
||||
groups. For commands and model support, see the
|
||||
[model-level user guide](../../../user_guide/diffusion/offloader/cpu_offload.md).
|
||||
|
||||
## Generic hook path
|
||||
|
||||
`ModelLevelOffloadBackend.enable()` discovers DiTs, encoders, and VAEs. It
|
||||
makes encoders and VAEs device resident for initialization, then installs
|
||||
`SequentialOffloadHook` instances:
|
||||
|
||||
- each DiT hook offloads all encoders and other DiTs before loading itself;
|
||||
- each encoder hook offloads every DiT before loading itself; and
|
||||
- VAEs are not part of the mutual-exclusion hook set.
|
||||
|
||||
The hook moves parameters and buffers without recursively calling
|
||||
`module.to()`. This avoids recursion through modules that retain references to
|
||||
other transformer blocks. Host tensors are pinned when configured; DTensors
|
||||
are not pinned through the ordinary tensor path.
|
||||
|
||||
Before forward returns control to the model, the hook synchronizes the active
|
||||
platform so compute cannot race the transfer. XPU uses blocking device-to-host
|
||||
copies because its allocator cannot safely evict storage while an asynchronous
|
||||
copy is pending.
|
||||
|
||||
## Split-model delegation
|
||||
|
||||
If a pipeline provides callable `enable_omni_model_cpu_offload`, the backend
|
||||
delegates instead of installing generic encoder/DiT hooks. This supports models
|
||||
whose mutually exclusive components are nested inside one transformer, such
|
||||
as a reasoner and generator.
|
||||
|
||||
The delegated pipeline owns its internal contexts, but must preserve the same
|
||||
contract:
|
||||
|
||||
- only the executing component is device resident;
|
||||
- enable and disable are idempotent;
|
||||
- transfers use the supplied device, pinning, and HSDP settings; and
|
||||
- disabling removes all model-local offload hooks or contexts.
|
||||
|
||||
The delegation is currently duck typed. If another independent implementation
|
||||
is added, formalize the enable/disable pair as a protocol rather than expanding
|
||||
attribute-name heuristics.
|
||||
|
||||
## Extension invariants
|
||||
|
||||
- New ordinary pipelines declare both DiT and encoder groups through
|
||||
`SupportsComponentDiscovery`.
|
||||
- Models with multiple DiTs offload the inactive DiTs as well as encoders.
|
||||
- Empty or parameterless modules do not trigger transfers.
|
||||
- Hook removal does not move modules; callers must not assume original
|
||||
residency is restored.
|
||||
|
||||
Shared selection and lifecycle behavior is defined in the
|
||||
[CPU Offloading design](README.md).
|
||||
@@ -0,0 +1,333 @@
|
||||
# Distributed Layerwise Offload
|
||||
|
||||
This document describes distributed layerwise offload (DLO) for diffusion
|
||||
models. DLO keeps only a small number of DiT blocks on the accelerator and
|
||||
streams the remaining blocks from host memory. The distributed backend can
|
||||
either shard those host-side weights across an existing parallel group or keep
|
||||
complete rank-local block sources and avoid an additional collective.
|
||||
|
||||
For user-facing commands, see the
|
||||
[distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
|
||||
and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
|
||||
|
||||
## Feature compatibility
|
||||
|
||||
Host-storage optimization and runtime compatibility are separate decisions.
|
||||
When direct checkpoint mmap is unavailable, DLO can still use tensors produced
|
||||
by the ordinary loader. "Compatibility path" below means that fallback is
|
||||
implemented but has less end-to-end coverage than the primary path.
|
||||
|
||||
Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsupported.
|
||||
|
||||
| Feature | DLO + AllGather | DLO without AllGather |
|
||||
|---|---|---|
|
||||
| **DP** | ✅ Primary path; host weights are sharded across the DP group. | ✅ Each DP rank streams complete rank-local blocks. |
|
||||
| **SP** | ✅ When DP=1, DLO uses the SP group for weight sharding. | ✅ SP remains active without a DLO weight collective. |
|
||||
| **TP > 1** | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. |
|
||||
| **HSDP** | ❌ Rejected to avoid double-sharding parameters. | ⚠️ Limited end-to-end coverage. |
|
||||
| **Per-tensor online FP8 linears** | ✅ Ordinary loader finalizes weights and scales before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
|
||||
| **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. |
|
||||
| **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. |
|
||||
| **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. |
|
||||
|
||||
See [Parallelism compatibility](#parallelism-compatibility) and
|
||||
[Request and loading constraints](#request-and-loading-constraints) for the
|
||||
detailed contracts and validation boundaries.
|
||||
|
||||
## Status
|
||||
|
||||
DLO is implemented for multi-device diffusion execution. The default
|
||||
AllGather path is the primary path for DP and SP deployments. The
|
||||
`--dlo-no-use-allgather` path streams complete blocks independently and adds no
|
||||
DLO weight collective.
|
||||
|
||||
Host storage is selected separately from the transfer protocol. The loader can
|
||||
produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
|
||||
otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on
|
||||
the same node can share immutable checkpoint pages when direct mmap is
|
||||
selected, while the ordinary-loader fallback still keeps a private runtime
|
||||
copy per process.
|
||||
|
||||
The Phase A shared-mmap support boundary is TP1. TP greater than one is an
|
||||
ordinary-loader compatibility path: DLO can consume the resulting TP-local
|
||||
tensors, but those configurations do not use checkpoint mmap and must not be
|
||||
used to claim shared-mmap host-memory savings.
|
||||
|
||||
The compatibility matrix below describes the current implementation. The
|
||||
unit-level guards are covered, but not every parallelism combination has a
|
||||
full model-and-hardware end-to-end test.
|
||||
|
||||
## Design
|
||||
|
||||
### DLO consumes the existing parallel topology
|
||||
|
||||
DLO does not create a new DP, TP, or SP topology. It reads the configured
|
||||
`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the
|
||||
standard distributed groups have been initialized.
|
||||
|
||||
The DLO weight-sharding group is selected as follows:
|
||||
|
||||
1. Use the existing DP group when `data_parallel_size > 1`.
|
||||
2. When DP is one and SP is greater than one, use the SP group.
|
||||
3. Otherwise, run rank-locally without a DLO process group.
|
||||
|
||||
TP is deliberately not used as DLO's AllGather group. HSDP has its own
|
||||
parameter-sharding lifecycle and is not allowed to be sharded a second time by
|
||||
DLO's AllGather path.
|
||||
|
||||
### The loader owns host-weight planning
|
||||
|
||||
Before it decides whether ordinary weight materialization can be skipped, the
|
||||
diffusion loader builds one `HostWeightPlan`. A direct-checkpoint mmap plan is
|
||||
accepted only when preflight proves all of the following:
|
||||
|
||||
- every required DiT parameter and persistent buffer has exactly one source;
|
||||
- runtime names, checkpoint keys, shapes, and dtypes match;
|
||||
- the runtime topology is TP1 without HSDP or online quantization; and
|
||||
- every custom loader operation is represented by a loader-owned checkpoint
|
||||
adapter.
|
||||
|
||||
The exact plan object is handed to DLO. The backend does not rescan checkpoint
|
||||
files, repeat the capability decision, or reconstruct names from its block
|
||||
topology. If preflight fails, the loader materializes weights normally and DLO
|
||||
consumes those runtime tensors.
|
||||
|
||||
The plan owns only dedicated DiT component sources. If a pipeline also exposes
|
||||
ordinary sources for a text encoder or another non-DiT component, the loader
|
||||
still consumes those sources and includes their loaded names in its strict
|
||||
coverage check. Only the source prefixes covered by the plan skip ordinary
|
||||
materialization. A source that mixes DiT and non-DiT weights fails closed to
|
||||
the complete ordinary-loader path because it cannot be skipped safely as a
|
||||
unit.
|
||||
|
||||
This boundary keeps checkpoint semantics out of DLO and avoids model-pipeline
|
||||
flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
|
||||
transforms. Model-specific direct-layout knowledge, when required, lives in a
|
||||
checkpoint adapter beside the ordinary loader.
|
||||
|
||||
### Phase-I normalized online-FP8 cache
|
||||
|
||||
The direct-checkpoint plan above cannot quantize a BF16 source online. Phase I
|
||||
adds an explicit, opt-in normalized cache for the narrower case of one TP1
|
||||
transformer with per-tensor online FP8 and DLO AllGather:
|
||||
|
||||
```text
|
||||
cache miss: BF16 source -> ordinary online FP8 loader -> atomic FP8 cache
|
||||
cache hit: FP8 cache -> direct mmap plan -> DLO host shard -> AllGather
|
||||
```
|
||||
|
||||
The cache stores the serialized runtime representation, including finalized FP8
|
||||
weights, generated scales, and persistent transformer buffers. Its manifest
|
||||
records the schema, component, quantization contract, tensor count, source
|
||||
files, and a source fingerprint. A missing, stale, incomplete, incompatible, or
|
||||
invalid artifact is a cache miss; cache publication is best-effort and falls
|
||||
back to the ordinary loader on failure.
|
||||
|
||||
Phase I keeps the cache binding deliberately simple: one source per transformer
|
||||
component, with existing per-tensor transforms applied before publication. The
|
||||
cache is not selected for DLO-disabled, no-AllGather, TP>1, or HSDP deployments,
|
||||
and it normalizes the transformer only. Cache sharing and compatibility for
|
||||
other components, transformed TP coordinates, and broader DP/SP layouts remain
|
||||
in [RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
|
||||
|
||||
### AllGather path
|
||||
|
||||
With the default `dlo_use_allgather=True`, each rank stores approximately
|
||||
`1 / group_size` of each streamable block in pinned host memory. The next
|
||||
block's shard is copied to a device buffer and reconstructed with
|
||||
`all_gather_into_tensor` on a communication stream while the current block is
|
||||
executing.
|
||||
|
||||
```text
|
||||
Compute: [Block N] [Block N+1] [Block N+2]
|
||||
H2D: [shard N+1] [shard N+2]
|
||||
AllGather: [full N+1] [full N+2]
|
||||
Buffers: [current slot] [prefetch slot] [current slot]
|
||||
```
|
||||
|
||||

|
||||
|
||||
The backend uses two shared device buffers, so accelerator weight residency is
|
||||
bounded by the largest streamed blocks rather than the complete model.
|
||||
|
||||
When direct checkpoint mmap is selected, the checkpoint mappings are only the
|
||||
source used to prepare each rank's persistent shard. They can be closed after
|
||||
shard preparation. Across the AllGather group, those private shards total
|
||||
approximately one runtime model copy.
|
||||
|
||||
An effective DLO group size of one performs no collective, even when
|
||||
`dlo_use_allgather=True`; it follows the rank-local transfer path described
|
||||
below.
|
||||
|
||||
When DP is greater than one, the engine can process one request per DP rank in
|
||||
the same denoising wave. Because AllGather is a collective, all participating
|
||||
requests must take the same execution path at every denoising step.
|
||||
|
||||
### Rank-local path without DLO AllGather
|
||||
|
||||
With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
|
||||
one and streams complete blocks using H2D copies only. The host backing may be
|
||||
either a loader-approved checkpoint mapping or ordinary runtime tensors.
|
||||
|
||||
For direct mmap, each process retains immutable safetensors views and uses two
|
||||
bounded pinned host staging slots. Processes on the same node that map the same
|
||||
files share physical checkpoint pages through the OS page cache. This removes
|
||||
the persistent private full-model copy per pure-DP process, but each process
|
||||
still packs and transfers every complete block. Sharing is node-local; each
|
||||
node has its own page cache.
|
||||
|
||||
When direct mmap preflight fails, the regular model loader remains responsible
|
||||
for preparing each rank's weights, including TP-local tensors or HSDP-managed
|
||||
parameters. In that fallback, each pure-DP process keeps a private full runtime
|
||||
copy.
|
||||
|
||||
This mode means:
|
||||
|
||||
- DP still provides independent replicas, but DLO does not shard weights
|
||||
across DP ranks.
|
||||
- SP still performs its normal activation/attention collectives, but DLO does
|
||||
not shard weights across SP ranks.
|
||||
- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
|
||||
DLO's additional weight AllGather is disabled.
|
||||
- Pure DP deployments share one checkpoint-backed copy per node when direct
|
||||
mmap is selected; the ordinary-loader fallback keeps one private runtime
|
||||
copy per rank.
|
||||
- The scheduler does not require a synchronized DP request wave for DLO.
|
||||
|
||||
## Parallelism compatibility
|
||||
|
||||
| Parallelism | DLO + AllGather | DLO without AllGather |
|
||||
|---|---|---|
|
||||
| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. |
|
||||
| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
|
||||
| **TP > 1** | Outside the Phase A shared-mmap support scope. The loader falls back before mutation, preserves TP-local layouts, and DLO may apply DP/SP host sharding to those ordinary runtime tensors. | Outside the Phase A shared-mmap support scope. The ordinary TP-aware loader produces rank-local tensors, which DLO streams without an additional weight collective; DP replicas retain private runtime storage. |
|
||||
| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
|
||||
|
||||
### Combined dimensions
|
||||
|
||||
- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater
|
||||
than one; SP continues to use its own sequence-parallel group. If DP is one,
|
||||
the SP group becomes DLO's sharding group in AllGather mode.
|
||||
- **DP + TP/SP without AllGather:** standard model loading defines the
|
||||
rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight
|
||||
collective.
|
||||
- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
|
||||
DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
|
||||
materialization and synchronization.
|
||||
- **HSDP + DP or TP:** rejected independently by the diffusion parallel
|
||||
configuration.
|
||||
|
||||
## Request and loading constraints
|
||||
|
||||
AllGather DP multi-concurrency requires:
|
||||
|
||||
- explicit `num_inference_steps`;
|
||||
- the same `num_inference_steps` for all requests in a wave; and
|
||||
- identical request arguments that affect the collective execution path.
|
||||
|
||||
The no-AllGather path does not impose these DLO-specific synchronized-wave
|
||||
requirements.
|
||||
|
||||
Direct checkpoint mmap can back either transfer path. It is currently limited
|
||||
to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
|
||||
ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the
|
||||
ordinary loader finalizes their runtime weights and scales; DLO then shards and
|
||||
reconstructs those tensors with their recorded layouts. Other online methods
|
||||
must use `--dlo-no-use-allgather` or disable online quantization until their
|
||||
runtime layouts are validated.
|
||||
|
||||
The Phase-I normalized online-FP8 cache is the supported TP1 transformer case.
|
||||
The general cache-compatibility contract for additional components, transformed
|
||||
TP coordinates, and cross-DP/SP reuse remains in
|
||||
[RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
|
||||
|
||||
## Validation coverage
|
||||
|
||||
Current source-level validation includes:
|
||||
|
||||
- HSDP + DLO + AllGather rejection;
|
||||
- HSDP + DLO without AllGather acceptance at configuration level;
|
||||
- loader preflight fallback for TP, HSDP, online quantization, unknown custom
|
||||
loaders, missing keys, and shape/dtype mismatches;
|
||||
- ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
|
||||
sharding of finalized weights and scales;
|
||||
- exact loader-to-backend plan transfer and ordinary-loader fallback;
|
||||
- rank-local mmap source retention, bounded two-slot staging, and adapter
|
||||
transforms without parameter-side flags;
|
||||
- resident-layer requests requiring no-AllGather;
|
||||
- DP request-wave validation for denoising-step compatibility;
|
||||
- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
|
||||
tests.
|
||||
|
||||
### B300 parallel-topology smoke matrix
|
||||
|
||||
A four-GPU B300 smoke test covered MiniMax-H3 FL2VA with the same prompt, seed,
|
||||
CUDNN attention backend, 256x256 output, two denoising steps, and
|
||||
`dlo_resident_layers=0`. The TP2 rows used DiT DP2xTP2 with the text encoder and
|
||||
VAEs at TP1. They validate the ordinary-loader fallback only, not direct mmap
|
||||
or shared-mmap host-memory savings.
|
||||
|
||||
| Configuration | Result | Warm E2E | Peak device memory | Host PSS |
|
||||
|---|---:|---:|---:|---:|
|
||||
| DP4xTP1 AllGather | Passed, 4 concurrent requests | 2.87 s / 4 requests | 13.84 GiB | 211.99 GiB |
|
||||
| DP4xTP1 no-AllGather | Passed, 1 request | 15.02 s | 13.23 GiB | 187.77 GiB |
|
||||
| DP2xTP2 AllGather | Passed, 2 concurrent requests | 4.16 s / 2 requests | 12.50 GiB | 211.97 GiB |
|
||||
| DP2xTP2 no-AllGather | Passed, 1 request | 3.51 s | 11.88 GiB | 314.01 GiB |
|
||||
|
||||
Within each topology, the AllGather and no-AllGather video and audio outputs
|
||||
were byte-identical. All four runs completed without an `ERROR` or traceback
|
||||
and released their device allocations. For DP4xTP1, no-AllGather direct mmap
|
||||
reduced total PSS by 24.22 GiB (11.4%) and `Private_Dirty` from 211.33 to
|
||||
125.32 GiB (40.7%) relative to AllGather. For DP2xTP2, preflight selected the
|
||||
ordinary loader as designed; no-AllGather PSS was 314.01 GiB, about 48% above
|
||||
AllGather, because DP replicas did not share checkpoint-backed runtime
|
||||
weights. This is a functional and memory smoke test, not a production-quality
|
||||
performance or output-quality benchmark.
|
||||
|
||||
### Host-memory measurement
|
||||
|
||||
A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the
|
||||
ordinary-loader fallback with direct mmap. Both runs used
|
||||
DP=2, TP=1, no DLO AllGather, BF16 weights, two denoising steps, and a
|
||||
256x256 four-second request. The ordinary-loader workers were sampled after
|
||||
initialization. The mmap workers were sampled after one completed request, so
|
||||
the checkpoint working set had been faulted into the page cache; this is the
|
||||
more conservative point for mmap.
|
||||
|
||||
The values below come from `/proc/<worker>/smaps_rollup` and include the whole
|
||||
worker, not only the DiT. The stable rank-to-rank difference comes from other
|
||||
pipeline components, so each worker should be compared with the same worker in
|
||||
the other storage mode.
|
||||
|
||||
| Worker | Ordinary RSS | mmap RSS | Ordinary PSS | mmap PSS | PSS reduction |
|
||||
|---|---:|---:|---:|---:|---:|
|
||||
| DP worker 0 | 168.27 GiB | 132.76 GiB | 167.84 GiB | 101.43 GiB | 66.40 GiB |
|
||||
| DP worker 1 | 116.19 GiB | 79.97 GiB | 115.73 GiB | 48.64 GiB | 67.09 GiB |
|
||||
| **Two-worker total** | — | — | **283.56 GiB** | **150.08 GiB** | **133.48 GiB (47.1%)** |
|
||||
|
||||
The direct-mmap workers each reported 62.45 GiB `Shared_Clean` but only
|
||||
31.20 GiB `Pss_File`, which is the proportional charge expected when the same
|
||||
resident checkpoint pages are mapped by two workers. `Private_Dirty` also fell
|
||||
from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker
|
||||
1, a reduction of about 97–98 GiB per worker. RSS understates this benefit
|
||||
because it counts a shared physical page in every process that maps it; summed
|
||||
PSS is the appropriate node-memory comparison.
|
||||
|
||||
The highest-value missing coverage is broader end-to-end numerical and
|
||||
lifecycle comparison against ordinary layerwise offload for DP+SP,
|
||||
HSDP+SP+no-AllGather, and TP greater than one across additional models and
|
||||
target CUDA/NCCL or CANN/HCCL hardware. That broader TP coverage does not
|
||||
change the Phase A direct-mmap TP1 support boundary.
|
||||
|
||||
## Recommendations
|
||||
|
||||
- Use **DP + DLO AllGather** for the supported throughput and host-memory
|
||||
scaling path.
|
||||
- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
|
||||
not the goal.
|
||||
- Use **no-AllGather** when independent replica execution is required. TP1
|
||||
direct-mmap deployments can share checkpoint pages per node; other layouts
|
||||
retain the ordinary loader's private host memory behavior and are outside
|
||||
the Phase A shared-mmap support scope.
|
||||
- Prefer **HSDP alone** for production HSDP deployments until the combined
|
||||
HSDP + DLO no-AllGather path has broader end-to-end coverage.
|
||||
@@ -0,0 +1,57 @@
|
||||
# Layerwise Offload
|
||||
|
||||
Layerwise offload streams transformer blocks through one accelerator. For
|
||||
commands and model support, see the
|
||||
[layerwise user guide](../../../user_guide/diffusion/offloader/layerwise_offload.md).
|
||||
|
||||
## Hook ring
|
||||
|
||||
`LayerWiseOffloadBackend` discovers every streamable block and registers one
|
||||
`LayerwiseOffloadHook` per block. Hooks form a ring: each hook owns the pinned
|
||||
host representation of its registered block and a reference to the next block
|
||||
it must prefetch; the final hook prefetches the first block for the next
|
||||
denoising iteration.
|
||||
|
||||
During initialization, parameters and buffers are grouped by dtype and copied
|
||||
into contiguous pinned CPU tensors. Original tensor storage is replaced with
|
||||
an empty or meta placeholder while identity and module structure are retained.
|
||||
Metadata records names, offsets, shapes, and dtypes for rematerialization.
|
||||
|
||||
## Forward lifecycle
|
||||
|
||||
For each block:
|
||||
|
||||
1. the previous hook's completion event ensures this block is materialized;
|
||||
2. the copy stream asynchronously allocates and fills the next block's device
|
||||
tensors;
|
||||
3. parameter and buffer storage is rebound to views of those tensors;
|
||||
4. the current block computes on the normal stream; and
|
||||
5. its post-forward hook replaces device storage with placeholders after all
|
||||
dependent work is safe.
|
||||
|
||||
The copy stream waits on the compute stream before reusing storage, and the
|
||||
compute stream waits on the per-block prefetch event before forward. These
|
||||
dependencies, not incidental global synchronization, define correctness.
|
||||
|
||||
## Topology contract
|
||||
|
||||
A DiT declares ordered block-container attributes using
|
||||
`_layerwise_offload_blocks_attrs` or an `OffloadPlan`. Each resolved item must
|
||||
be an executable `nn.Module`; ordering must match forward execution. Multiple
|
||||
containers are concatenated in execution order.
|
||||
|
||||
Non-block DiT modules, encoders, VAEs, and declared resident modules are not
|
||||
streamed by this backend.
|
||||
|
||||
## Invariants and limitations
|
||||
|
||||
- Each offloaded tensor has one pinned host backing store.
|
||||
- Parameter identity is preserved when storage is rebound.
|
||||
- Blocks with heterogeneous shapes and dtypes use their recorded metadata;
|
||||
implementations must not assume identical blocks.
|
||||
- The ring is initialized before the first forward so block zero is available.
|
||||
- This backend is single-device; distributed weight reconstruction belongs to
|
||||
[Distributed Layerwise Offload](distributed_layerwise_offload.md).
|
||||
|
||||
Shared selection and lifecycle behavior is defined in the
|
||||
[CPU Offloading design](README.md).
|
||||
@@ -76,6 +76,12 @@ equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter.
|
||||
ModelOpt checkpoint metadata is detected from its method, producer metadata, or
|
||||
`quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration.
|
||||
|
||||
Unless the input is a per-component dictionary, the resolved configuration is
|
||||
passed unchanged to every quantization-aware component constructed by the
|
||||
pipeline. It affects only layers supported by that component's quantization
|
||||
method; it does not rewrite arbitrary `torch.nn` modules. A per-component
|
||||
dictionary narrows this global scope through runtime layer-prefix routing.
|
||||
|
||||
The factory uses Omni-specific overrides for methods whose diffusion behavior
|
||||
or platform dispatch is not provided by the vLLM registry. The current override
|
||||
set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/
|
||||
@@ -154,9 +160,11 @@ quantizable vLLM layers. Each method configuration is responsible for:
|
||||
Model code must keep precision-sensitive or unsupported modules unquantized.
|
||||
The shared `safe_quant_config()` helper is used by model integrations for
|
||||
normalization and modulation layers where a generic FP8 configuration would be
|
||||
unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers
|
||||
are not implicitly quantized just because the main transformer is quantized;
|
||||
they require explicit component routing and method support.
|
||||
unsafe. Encoders exposed through the quantization factory receive a global
|
||||
configuration just like the main transformer, but only their supported vLLM
|
||||
quantizable layers are affected. Ordinary `torch.nn` encoder layers, VAEs,
|
||||
schedulers, and tokenizers are not rewritten automatically. A component map can
|
||||
select a narrower scope, but it cannot add quantization support to a layer.
|
||||
|
||||
Weight quantization and quantized KV cache are separate paths. The former is
|
||||
resolved through `quantization_config`; runtime attention/KV-cache quantization
|
||||
@@ -182,12 +190,15 @@ configured parallelism strategy:
|
||||
- `WeightsMapper` transformations must be applied before matching ignored-layer
|
||||
lists or component prefixes.
|
||||
|
||||
Distributed layerwise offload has an additional boundary: its sharded mmap
|
||||
AllGather path is not compatible with online quantization because the path
|
||||
expects already materialized checkpoint weights. Use the rank-local
|
||||
`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when
|
||||
combining these features. See the [DLO design document](distributed_layerwise_offload.md)
|
||||
for the full compatibility matrix.
|
||||
Distributed layerwise offload has an additional boundary: direct checkpoint
|
||||
mmap cannot prepare online-quantized tensors because it bypasses the ordinary
|
||||
loader that creates their runtime weights and scales. Per-tensor online FP8
|
||||
linears can still use DLO AllGather through the ordinary loader: each rank first
|
||||
finalizes the FP8 tensors, then DLO shards those runtime tensors for
|
||||
reconstruction during prefetch. Other online methods remain limited to
|
||||
`--dlo-no-use-allgather` until their packing and scale layouts are validated.
|
||||
See the [DLO design document](offloader/distributed_layerwise_offload.md) for
|
||||
the full compatibility matrix.
|
||||
|
||||
## Adding a quantization backend
|
||||
|
||||
|
||||
@@ -0,0 +1,213 @@
|
||||
# Realtime AR-Diffusion sessions
|
||||
|
||||
> **Status:** experimental. This document describes the internal execution
|
||||
> contract implemented under `vllm_omni.experimental.ar_diffusion`. It does not
|
||||
> define a public HTTP or WebSocket schema.
|
||||
|
||||
## Scope
|
||||
|
||||
Realtime AR-Diffusion serves world models that generate one autoregressive
|
||||
latent block per request while retaining model state across requests. The
|
||||
generic runtime owns identity, event ordering, session lifecycle, routing, and
|
||||
KV allocation. A model adapter owns model-specific controls and conditioning.
|
||||
|
||||
LingBot World v2 is the first single-KV-branch integration. DreamZero uses the
|
||||
same runner capability with two logical CFG branches.
|
||||
|
||||
## Identities
|
||||
|
||||
Four identifiers have distinct lifetimes:
|
||||
|
||||
| Identifier | Lifetime | Owner | Contract |
|
||||
| --- | --- | --- | --- |
|
||||
| `session_id` | One persistent world | Session manager | Selects runner KV and model-owned state. |
|
||||
| `event_id` | One control/prompt update | Transport adapter | Unique and monotonically increasing within a session, including across reset. |
|
||||
| `chunk_index` | One committed AR block position | Session | Contiguous from zero; reset starts again at zero. |
|
||||
| `request_id` | One chunk execution | Session | Correlates the tick with its returned metadata. |
|
||||
|
||||
`AsyncOmni` keeps its normal UUID-suffixed engine routing ID. The tick
|
||||
`request_id` remains inside the immutable AR-Diffusion snapshot and the
|
||||
standard output metadata, while caller-visible `OmniRequestOutput.request_id`
|
||||
continues to use AsyncOmni's existing external-ID contract. Protocol identity
|
||||
therefore does not require a private entrypoint override.
|
||||
|
||||
## Request lifecycle
|
||||
|
||||
1. `ARDiffusionSession.accept_event()` validates and queues an immutable
|
||||
`ARDiffusionSessionEvent`. Prompt and opaque control tracks may be updated in
|
||||
the same event.
|
||||
2. `ARDiffusionSession.next_chunk()` serializes chunk execution and calls
|
||||
`_build_tick_locked()`. The resulting `ARDiffusionTickRequest` is a snapshot
|
||||
of the current prompt, reduced controls, ordered event IDs, and chunk index.
|
||||
3. `ARDiffusionOmniTickConsumer.execute_tick()` copies the stage sampling
|
||||
parameters, stores the typed tick in
|
||||
`sampling_params.extra_args["ar_diffusion_tick"]`, constructs the normal Omni
|
||||
prompt through its `prompt_provider`, and invokes `AsyncOmni.generate()`.
|
||||
4. `AsyncOmni` allocates its normal unique engine routing ID and submits one
|
||||
`OmniDiffusionRequest` through the orchestrator.
|
||||
5. `ARDiffusionModelRunner._request_session()` parses the typed tick.
|
||||
`_get_or_create_session()` obtains runner-owned paged KV, and
|
||||
`bind_ar_diffusion_state()` exposes it to the pipeline only for the duration
|
||||
of `execute_model()`.
|
||||
6. The model pipeline validates the tick and generates exactly one latent
|
||||
block. LingBot interprets camera controls in
|
||||
`LingBotCameraControlReducer` and
|
||||
`LingBotWorldCausalDMDPipeline.forward()`; the generic runtime does not parse
|
||||
camera or action schemas.
|
||||
7. The pipeline returns one `DiffusionOutput` with the generated payload and an
|
||||
`ARDiffusionChunkMetadata` mapping under `metadata["ar_diffusion"]`.
|
||||
`AsyncOmni` exposes it as
|
||||
`OmniRequestOutput.multimodal_output["metadata"]["ar_diffusion"]`.
|
||||
8. The consumer parses the standard envelope. The session commits its reducer
|
||||
state, prompt, controls, event IDs, and next chunk index only when returned
|
||||
metadata exactly equals the submitted tick snapshot.
|
||||
|
||||
The metadata envelope is:
|
||||
|
||||
```json
|
||||
{
|
||||
"metadata": {
|
||||
"ar_diffusion": {
|
||||
"session_id": "world-7",
|
||||
"request_id": "ar-world-7-3-...",
|
||||
"chunk_index": 3,
|
||||
"applied_event_ids": [10, 11]
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Transaction and failure semantics
|
||||
|
||||
Accepted events are not acknowledged as applied until the model output and
|
||||
metadata are validated. Reducer `prepare()` is speculative; reducer `commit()`
|
||||
and the session snapshot form one logical commit. A model, metadata, or reducer
|
||||
failure leaves the chunk index and queued events uncommitted and transitions the
|
||||
session to `FAILED`.
|
||||
|
||||
After a tick failure, worker state is closed because a forward may have
|
||||
partially modified or already released KV. The same chunk is not retried in
|
||||
place. A caller must explicitly reset the session, starting again at chunk
|
||||
zero, or close it.
|
||||
|
||||
Explicit close follows this state machine:
|
||||
|
||||
```text
|
||||
ACTIVE or FAILED
|
||||
|
|
||||
v
|
||||
CLOSING ---- worker cleanup succeeds ----> CLOSED
|
||||
|
|
||||
+------- worker cleanup fails --------> CLEANUP_FAILED
|
||||
|
|
||||
+-- close retry --> CLOSING
|
||||
```
|
||||
|
||||
`CLEANUP_FAILED` is a local tombstone. The session manager retains it, rejects
|
||||
creation of the same `session_id`, and rejects reset. Only an idempotent close
|
||||
retry that succeeds removes the session. Transport disconnect uses the same
|
||||
path.
|
||||
|
||||
## KV ownership and capacity
|
||||
|
||||
`ARDiffusionKVCacheSpec` is the pipeline-owned immutable geometry contract. It
|
||||
declares:
|
||||
|
||||
- TP-local self-attention geometry, AR block size, sink, and recent window;
|
||||
- logical KV branches and their worker-local slots;
|
||||
- fixed-length cross-attention caches;
|
||||
- scratch tokens required by an uncommitted forward;
|
||||
- persistent model-owned CUDA bytes per resident session; and
|
||||
- the requested resident-session capacity.
|
||||
|
||||
The runner computes:
|
||||
|
||||
```text
|
||||
required(capacity) =
|
||||
scratch
|
||||
+ managed_self_attention_pages(capacity)
|
||||
+ capacity * cross_attention_bytes_per_session
|
||||
+ capacity * model_owned_state_bytes_per_session
|
||||
```
|
||||
|
||||
`gpu_memory_fraction` is a soft expansion budget:
|
||||
|
||||
```text
|
||||
configured_budget = available_device_bytes * gpu_memory_fraction
|
||||
effective_budget = max(configured_budget, required(1))
|
||||
```
|
||||
|
||||
If `required(1)` exceeds actual available device memory, initialization fails.
|
||||
Otherwise the runner always admits one viable session and uses the configured
|
||||
fraction to select additional resident sessions up to the model-declared cap.
|
||||
All model-owned reservations are deducted before the paged self-attention pool
|
||||
is allocated.
|
||||
|
||||
LingBot reports its persistent full-horizon image-condition tensor. DreamZero
|
||||
reports the measured 603 MiB per-session Wan VAE causal-convolution state.
|
||||
|
||||
## Concurrency and routing limits
|
||||
|
||||
The current runner executes one request at a time:
|
||||
|
||||
- `max_num_seqs` must be one;
|
||||
- request batching and diffusion step execution are rejected;
|
||||
- one AR block is generated per engine request; and
|
||||
- an AR-Diffusion stage must have exactly one replica.
|
||||
|
||||
The single-replica restriction is intentional. Multi-replica support requires
|
||||
session-affine routing so every tick, reset, and close reaches the worker that
|
||||
owns the session. Random or round-robin routing is not correct.
|
||||
|
||||
Multiple resident sessions may share one runner. Capacity selection accounts
|
||||
for all persistent pools, and the runner performs LRU eviction through the same
|
||||
model lifecycle hook used by explicit close.
|
||||
|
||||
## Model boundary
|
||||
|
||||
The generic runtime transports controls as opaque
|
||||
`ARDiffusionControlInput(track, schema, data)` values. Model adapters validate
|
||||
and reduce those values. This keeps target/velocity, key state, pose,
|
||||
intrinsics, and other model-specific conventions out of the session protocol.
|
||||
|
||||
For LingBot World v2:
|
||||
|
||||
- `lingbot_world/actions.py` parses key-state and trajectory controls;
|
||||
- `lingbot_world/camera.py` loads or constructs camera trajectories and Plücker embeddings;
|
||||
- `lingbot_world/transformer.py` implements checkpoint-compatible causal attention; and
|
||||
- `lingbot_world/pipeline.py` constructs conditioning, owns small non-KV session state, and
|
||||
produces one block plus the standard metadata envelope.
|
||||
|
||||
## Non-goals
|
||||
|
||||
This contract does not currently provide:
|
||||
|
||||
- a public HTTP or WebSocket request schema (the structured-interaction
|
||||
frontend is tracked separately in
|
||||
[#5527](https://github.com/vllm-project/vllm-omni/pull/5527));
|
||||
- camera/action semantics shared by every world model;
|
||||
- session migration or replication across workers;
|
||||
- retry of an ambiguous partially executed chunk;
|
||||
- cross-stage KV transfer; or
|
||||
- stateful streaming VAE decode.
|
||||
|
||||
Those features can be layered on top without adding LingBot-specific identity
|
||||
or lifecycle fields to the generic runtime.
|
||||
|
||||
## Required regression coverage
|
||||
|
||||
CPU contract tests cover:
|
||||
|
||||
- separation of tick correlation IDs from internal engine routing IDs;
|
||||
- event ordering, metadata equality, and reducer fail-closed commit;
|
||||
- close/disconnect failure tombstones and cleanup retry;
|
||||
- capacity one under the default fraction for shipped LingBot geometry;
|
||||
- capacity expansion under a tuned fraction;
|
||||
- DreamZero model-owned state accounting;
|
||||
- runner session reuse, reset, close, LRU eviction, and forward cleanup; and
|
||||
- LingBot model registration, imports, camera/action reduction, and one-block
|
||||
output metadata.
|
||||
|
||||
GPU validation must additionally exercise real model loading, action input,
|
||||
prompt switching, rolling-window replacement, reset, close, finite output
|
||||
latents, and exact metadata/request identity.
|
||||
@@ -2,8 +2,10 @@
|
||||
|
||||
Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. Usage — the config keys
|
||||
and how to pick an operating point — is in
|
||||
[Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). This page explains
|
||||
the algorithm.
|
||||
[TRTLLM Attention](../../user_guide/diffusion/attention_backends/trtllm.md#skip-softmax).
|
||||
The shared selector contract is documented in
|
||||
[Diffusion Attention Backend Selection](attention_backend_selection.md). This
|
||||
page explains the algorithm.
|
||||
|
||||
## Motivation
|
||||
|
||||
|
||||
+11
-3
@@ -22,6 +22,7 @@ implementation contract; it is not, by itself, a general support claim.
|
||||
- [Async Diffusion Output](feature/async_diffusion_output.md)
|
||||
- [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
|
||||
- [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
|
||||
- [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md)
|
||||
|
||||
### Communication
|
||||
|
||||
@@ -53,15 +54,22 @@ implementation contract; it is not, by itself, a general support claim.
|
||||
#### Attention optimization
|
||||
|
||||
The [Diffusion Attention Backends](../user_guide/diffusion/attention_backends.md)
|
||||
guide lists all selectable backends and their platform defaults. Skip-Softmax
|
||||
is the backend-related optimization with a standalone feature-design contract:
|
||||
guides list selectable backends, platform defaults, installation, and tuning.
|
||||
The design contracts separate selection mechanics from backend algorithms:
|
||||
|
||||
- [Attention Backend Selection](feature/attention_backend_selection.md)
|
||||
- [Skip-Softmax](feature/skip_softmax.md)
|
||||
|
||||
#### CPU offloading
|
||||
|
||||
- [Overview and Shared Contracts](feature/offloader/README.md)
|
||||
- [Model-Level Offload](feature/offloader/cpu_offload.md)
|
||||
- [Layerwise Offload](feature/offloader/layerwise_offload.md)
|
||||
- [Distributed Layerwise Offload](feature/offloader/distributed_layerwise_offload.md)
|
||||
|
||||
- [Cache-DiT](feature/cache_dit.md)
|
||||
- [TeaCache](feature/teacache.md)
|
||||
- [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
|
||||
- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
|
||||
|
||||
## Infrastructure and Performance
|
||||
|
||||
|
||||
@@ -18,8 +18,8 @@ validation_paths:
|
||||
- tests/config/**
|
||||
upstream_refs:
|
||||
- vllm.config
|
||||
last_reviewed: 2026-08-07
|
||||
last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
|
||||
last_reviewed: 2026-08-11
|
||||
last_verified_commit: e356708da4cd39992405686dacfc32b18fcccc7f
|
||||
---
|
||||
|
||||
# vLLM-Omni configuration
|
||||
@@ -39,11 +39,66 @@ The PR review proposal to capture environment-derived values in configuration
|
||||
objects at initialization is an open design question, not a current
|
||||
invariant.
|
||||
|
||||
## Current structured-stage reuse behavior
|
||||
|
||||
The current RFC #4021 implementation uses upstream vLLM config classes as the
|
||||
schema source for AR and generation stages:
|
||||
|
||||
- `OmniStageLoadConfig` inherits `vllm.config.LoadConfig`.
|
||||
- `OmniStageCacheConfig` inherits `vllm.config.CacheConfig`.
|
||||
- `OmniStageSchedulerConfig` inherits `vllm.config.SchedulerConfig`.
|
||||
- `OmniStageParallelConfig` inherits `vllm.config.ParallelConfig`.
|
||||
|
||||
This inheritance preserves the existing load, cache, scheduler, and parallel
|
||||
concern boundaries while reducing duplicate downstream declarations. It does
|
||||
not, by itself, make every inherited field an effective vLLM-Omni runtime
|
||||
option. Construction and runtime consumption are separate surfaces:
|
||||
|
||||
| Surface | Current behavior |
|
||||
| --- | --- |
|
||||
| Structured schema | Upstream dataclass fields are inherited. Their defaults, default factories, and applicable Pydantic validation can participate when the structured object is constructed. |
|
||||
| Omni input ownership | Pipeline/deploy/CLI construction continues to accept only fields with an existing structured owner. Direct construction of an inherited sub-config can expose a wider upstream schema. |
|
||||
| Engine projection | For AR and generation stages, reusable fields are discovered by intersecting each upstream config dataclass with upstream `EngineArgs`. Only constructor-explicit values are emitted; inherited defaults remain deferred to terminal materialization. |
|
||||
| Terminal materialization | The engine-owning process constructs the final upstream `VllmConfig` and performs model-, platform-, rank-, port-, and backend-dependent initialization. |
|
||||
|
||||
The generated projection maps include known upstream naming differences such
|
||||
as `cache_dtype` to `kv_cache_dtype`, `policy` to `scheduling_policy`, and
|
||||
`data_parallel_master_ip` to `data_parallel_address`. A field added to both an
|
||||
upstream concern config and `EngineArgs` therefore enters the AR/generation
|
||||
projection without requiring a second downstream allowlist. Explicit inherited
|
||||
fields that have no projection raise an error rather than being silently
|
||||
dropped.
|
||||
|
||||
Ownership exclusions remain deliberate. Stage topology owns `scheduler_cls`,
|
||||
cache owns `disable_hybrid_kv_cache_manager`, and Omni runtime owns
|
||||
`distributed_executor_backend` and `worker_cls`; vLLM's private API-process
|
||||
fields are terminal internals. These exclusions prevent one input from being
|
||||
constructed or projected through two config concerns. Effective-engine-argument
|
||||
tests cover both the dynamic mapping and this exclusion boundary.
|
||||
|
||||
`CompilationConfig` and `ProfilerConfig` are direct structured-stage fields:
|
||||
`stage.compilation_config` and `stage.profiler_config`. They accept mapping/YAML
|
||||
inputs and complete upstream value objects. Mapping inputs are validated and
|
||||
materialized as their concrete upstream types when the structured config is
|
||||
constructed, so downstream consumers see one resolved representation.
|
||||
Prebuilt upstream objects retain their type across the structured projection
|
||||
boundary.
|
||||
|
||||
Quantization remains on its existing Omni-owned transport contract. This
|
||||
reuse step does not adopt upstream `QuantizationConfigArgs`, because the
|
||||
current engine-specific quantization split needs to be resolved separately.
|
||||
|
||||
Diffusion keeps its existing engine-field ownership and projection surface.
|
||||
Shared Python inheritance must not be treated as evidence that an LLM-only
|
||||
parallel input is effective for a diffusion stage.
|
||||
|
||||
## Safe-change guide while deferred
|
||||
|
||||
Changes should trace every affected producer to its runtime consumer and test
|
||||
the structured, legacy, CLI, and deployment paths that are actually supported.
|
||||
Do not infer a stable contract from this placeholder page.
|
||||
For inherited vLLM fields, tests should separately cover constructor behavior
|
||||
and effective `EngineArgs` projection. Do not infer a stable contract from
|
||||
this placeholder page.
|
||||
|
||||
## Promotion gate
|
||||
|
||||
|
||||
+11
-5
@@ -20,13 +20,19 @@ an implementation contract has a user-facing workflow.
|
||||
| Choose serial, batched, step-wise, or streaming diffusion execution | [Execution Modes and Streaming](../user_guide/diffusion/execution_modes.md) | [Diffusion Continuous Batching](../design/feature/diffusion_continuous_batching.md), [Async Diffusion Output](../design/feature/async_diffusion_output.md) |
|
||||
| Reclaim stage memory without restarting the server | [Sleep Mode](sleep_mode.md) | Runtime lifecycle behavior is documented in the user guide |
|
||||
|
||||
The [Prefill-Decode Disaggregation design reference](pd_disaggregation.md)
|
||||
describes the experimental Qwen3-Omni runtime topology and its current
|
||||
integration boundaries. It is not yet a supported launch recipe.
|
||||
|
||||
Some runtime designs are deliberately not promoted as standalone User Guide
|
||||
features yet:
|
||||
|
||||
- [Disaggregated Inference](../design/feature/disaggregated_inference.md) and
|
||||
- Generic, OmniConnector-based
|
||||
[Disaggregated Inference](../design/feature/disaggregated_inference.md) is not
|
||||
yet a standalone User Guide feature. Its
|
||||
[OmniConnector implementations](../design/index.md#communication) describe
|
||||
topology and transport contracts. Practical deployment configuration remains
|
||||
under [Pipeline and deploy configurations](../configuration/stage_configs.md).
|
||||
transport contracts, while practical connector configuration remains under
|
||||
[Pipeline and deploy configurations](../configuration/stage_configs.md).
|
||||
- [Async Chunk](../design/feature/async_chunk.md) and
|
||||
[Async Omni Output Materialization](../design/feature/omni_async_output_materialization.md)
|
||||
are model- and pipeline-dependent. Use the selected model's deploy
|
||||
@@ -48,10 +54,10 @@ defines the shared configuration and backend extension points.
|
||||
| Goal | User guide | Related design contract |
|
||||
| --- | --- | --- |
|
||||
| Compare acceleration methods and supported combinations | [Overview](../user_guide/diffusion_features.md), [Feature Compatibility](../user_guide/feature_compatibility.md) | [Diffusion acceleration designs](../design/index.md#diffusion-acceleration) |
|
||||
| Move weights between host and device memory | [CPU Offloading](../user_guide/diffusion/cpu_offload.md) | [Distributed Layerwise Offload](../design/feature/distributed_layerwise_offload.md) |
|
||||
| Move weights between host and device memory | [CPU Offloading](../user_guide/diffusion/cpu_offload.md) | [CPU Offloading](../design/feature/offloader/README.md) |
|
||||
| Reuse denoising computation | [Cache Acceleration](../user_guide/diffusion/cache_acceleration/cache_dit.md) | [Cache-DiT](../design/feature/cache_dit.md), [TeaCache](../design/feature/teacache.md) |
|
||||
| Distribute diffusion work across devices | [Parallelism](../user_guide/diffusion/parallelism/overview.md) | [Parallelism designs](../design/index.md#parallelism) |
|
||||
| Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Skip-Softmax](../design/feature/skip_softmax.md) |
|
||||
| Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Attention Backend Selection](../design/feature/attention_backend_selection.md) |
|
||||
| Compile repeated diffusion regions | [Regional Compilation](../user_guide/diffusion/regional_compilation.md) | User-facing optimization guide |
|
||||
| Add generated video frames | [Frame Interpolation](../user_guide/diffusion/frame_interpolation.md) | User-facing extension guide |
|
||||
| Reduce diffusion model startup time | [Startup and Loading](../user_guide/diffusion/startup_and_loading.md) | User-facing loading guide |
|
||||
|
||||
@@ -101,7 +101,7 @@ outputs = omni.generate(
|
||||
)
|
||||
|
||||
# Access custom trajectory data
|
||||
output = outputs[0].request_output
|
||||
output = outputs[0]
|
||||
print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
|
||||
print(f"Trajectory latents shape: {output.latents.shape}")
|
||||
```
|
||||
|
||||
@@ -0,0 +1,140 @@
|
||||
# Prefill-Decode (PD) Disaggregation (experimental)
|
||||
|
||||
!!! warning "Experimental design reference"
|
||||
PD disaggregation is only partially integrated on the current `main` branch.
|
||||
The runtime contains PD detection and routing scaffolding, but there is no
|
||||
supported deploy overlay or end-to-end validated launch recipe yet. Treat this
|
||||
page as a description of the current design, not as production guidance.
|
||||
|
||||
Prefill-decode disaggregation splits the Qwen3-Omni Thinker into two logical
|
||||
stages: a prefill worker that processes the prompt and produces KV cache, and a
|
||||
decode worker that imports that cache and generates tokens. The Talker and
|
||||
Code2Wav stages remain downstream; PD does not disaggregate the entire
|
||||
Qwen3-Omni pipeline.
|
||||
|
||||
## Current design
|
||||
|
||||
### Runtime topology
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
client["Client request"] --> orchestrator["Omni Orchestrator"]
|
||||
orchestrator -->|"original prompt<br/>max_tokens = 1"| prefill["Thinker prefill<br/>KV producer"]
|
||||
prefill -->|"KV blocks"| connector[("vLLM KV connector")]
|
||||
prefill -->|"KV routing metadata"| orchestrator
|
||||
orchestrator -->|"original prompt<br/>remote KV parameters"| decode["Thinker decode<br/>KV consumer"]
|
||||
connector -->|"remote KV load"| decode
|
||||
decode -->|"generated text"| client
|
||||
```
|
||||
|
||||
The orchestrator owns the request lifecycle, while the vLLM KV connector owns
|
||||
the bulk KV transfer. The original prompt is submitted to both Thinker stages:
|
||||
the decode worker still needs the request tokens and multimodal features to
|
||||
construct its request, while the imported KV cache prevents it from repeating
|
||||
the completed prefill computation.
|
||||
|
||||
For the full Qwen3-Omni pipeline, splitting the Thinker adds one logical stage:
|
||||
|
||||
| Stage | Role | Output or handoff |
|
||||
| --- | --- | --- |
|
||||
| 0 | Thinker prefill | Saves prompt KV and prompt-side multimodal output |
|
||||
| 1 | Thinker decode | Loads remote KV, generates text, and produces decode-side conditioning |
|
||||
| 2 | Talker | Converts combined Thinker conditioning into codec codes |
|
||||
| 3 | Code2Wav | Converts codec codes into audio |
|
||||
|
||||
Text remains a final output of the decode stage, and audio remains a final
|
||||
output of Code2Wav.
|
||||
|
||||
### Request lifecycle
|
||||
|
||||
1. At startup, PD detection looks for one stage marked `is_prefill_only` and
|
||||
one downstream stage marked `is_decode_only`. The decode stage must list the
|
||||
prefill stage as an input source.
|
||||
2. The orchestrator clones the Thinker sampling parameters for prefill, forces
|
||||
`max_tokens=1`, clears stop conditions, and sets producer-side
|
||||
`kv_transfer_params`. This makes prefill finish by length after exporting KV.
|
||||
3. When prefill finishes, the orchestrator captures its KV routing metadata and
|
||||
Qwen multimodal output under the original request ID.
|
||||
4. The orchestrator submits the original prompt to decode with consumer-side
|
||||
`kv_transfer_params`. Those parameters identify the transfer, prefill engine,
|
||||
bootstrap endpoint, and, when required by the connector, remote request.
|
||||
5. The decode engine imports the KV cache and continues generation. The normal
|
||||
Omni stage path then carries Thinker conditioning through Talker and
|
||||
Code2Wav.
|
||||
|
||||
### Control and data paths
|
||||
|
||||
PD uses two different connector layers. They solve different handoffs and are
|
||||
not interchangeable.
|
||||
|
||||
| Path | Payload | Owner |
|
||||
| --- | --- | --- |
|
||||
| PD control path | Request ID, original prompt, `transfer_id`, bootstrap address, engine ID, and connector-specific request metadata | Omni Orchestrator |
|
||||
| PD data path | Thinker attention KV blocks | vLLM KV connector |
|
||||
| Downstream model path | Thinker hidden states and embeddings, Talker codec codes, and audio tensors | vLLM-Omni stage connectors and processors |
|
||||
|
||||
### Detection and validation
|
||||
|
||||
The current runtime accepts one logical prefill/decode pair per pipeline. Each
|
||||
stage may use tensor parallelism, but both sides of the pair must satisfy these
|
||||
invariants:
|
||||
|
||||
- both stages define `kv_transfer_config`
|
||||
- prefill uses `kv_producer` or `kv_both`
|
||||
- decode uses `kv_consumer` or `kv_both`
|
||||
- both stages use the same KV connector
|
||||
- `kv_buffer_device` and `kv_buffer_size` agree when set on both stages
|
||||
- both stages use the same `tensor_parallel_size`
|
||||
|
||||
These checks happen before request processing so an inconsistent PD topology
|
||||
fails during engine initialization.
|
||||
|
||||
## What is implemented today
|
||||
|
||||
| Area | Current state |
|
||||
| --- | --- |
|
||||
| Pair detection and topology validation | Implemented in the shared PD helper |
|
||||
| Prefill sampling preparation | Implemented for sync and async entrypoints |
|
||||
| Prefill-to-decode routing | Implemented in the current Orchestrator |
|
||||
| Deploy-based configuration | Not yet migrated; the supported `stages:` deploy schema cannot currently expand Qwen3-Omni into the four-stage PD topology |
|
||||
| Qwen3 Thinker-to-Talker handoff | Bridge state and embedding-merge helpers exist, but the merge is not wired into the live stage processor |
|
||||
| Connector compatibility | Incomplete; the Orchestrator currently requires `remote_request_id`, while the pinned Mooncake connector completes prefill without returning that metadata |
|
||||
| Validation coverage | PD entrypoint tests and the PD deploy overlay in online end-to-end tests are temporarily disabled during migration |
|
||||
| Performance | No non-regression or throughput/latency result is established for the current path |
|
||||
|
||||
Because these pieces are still converging, the previous legacy `stage_args` YAML
|
||||
and `vllm serve` command have been removed. The serve CLI now accepts the
|
||||
deploy-based `stages:` schema, so retaining the legacy recipe would imply a
|
||||
launch path that current `main` rejects.
|
||||
|
||||
## Configuration contract under development
|
||||
|
||||
The intended deploy integration starts from the
|
||||
[default Qwen3-Omni deploy config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
|
||||
and expands its frozen three-stage pipeline into the four stages shown above.
|
||||
A complete integration must:
|
||||
|
||||
- split the Thinker definition without duplicating the Talker or Code2Wav model
|
||||
contracts
|
||||
- remap downstream stage IDs, input sources, and stage connectors
|
||||
- inject producer and consumer `kv_transfer_config` values into the two Thinker
|
||||
stages
|
||||
- allocate separate device sets for prefill and decode
|
||||
- preserve both final outputs: decode text and Code2Wav audio
|
||||
- provide connector-specific routing metadata without hard-coding one connector
|
||||
contract
|
||||
|
||||
The required number of GPUs depends on tensor-parallel sizes and whether Talker
|
||||
and Code2Wav share a device. A basic TP=1 layout uses one device for prefill, one
|
||||
for decode, and one for the two downstream audio stages.
|
||||
|
||||
## Implementation map
|
||||
|
||||
- PD detection, validation, and sampling-parameter preparation:
|
||||
[`PDDisaggregationMixin`](gh-file:vllm_omni/entrypoints/pd_utils.py)
|
||||
- startup detection and connector metadata:
|
||||
[`AsyncOmniEngine`](gh-file:vllm_omni/engine/async_omni_engine.py)
|
||||
- request-scoped prefill-to-decode routing:
|
||||
[`Orchestrator`](gh-file:vllm_omni/engine/orchestrator.py)
|
||||
- Qwen3 downstream conditioning helpers:
|
||||
[Qwen3-Omni stage input processor](gh-file:vllm_omni/model_executor/stage_input_processors/qwen3_omni.py)
|
||||
@@ -47,7 +47,7 @@ if __name__ == "__main__":
|
||||
omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
|
||||
prompt = "a cup of coffee on the table"
|
||||
outputs = omni.generate(prompt)
|
||||
images = outputs[0].request_output.images
|
||||
images = outputs[0].images
|
||||
images[0].save("coffee.png")
|
||||
```
|
||||
|
||||
@@ -75,8 +75,7 @@ if __name__ == "__main__":
|
||||
]
|
||||
omni_outputs = omni.generate(prompts)
|
||||
for i_prompt, prompt_output in enumerate(omni_outputs):
|
||||
this_request_output = prompt_output.request_output
|
||||
this_images = this_request_output.images
|
||||
this_images = prompt_output.images
|
||||
for i_image, image in enumerate(this_images):
|
||||
image.save(f"p{i_prompt}-img{i_image}.jpg")
|
||||
print("saved to", f"p{i_prompt}-img{i_image}.jpg")
|
||||
|
||||
@@ -37,15 +37,23 @@ th {
|
||||
| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | — |
|
||||
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — |
|
||||
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | — |
|
||||
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
|
||||
| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | — |
|
||||
| `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — |
|
||||
| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | | — |
|
||||
| `LTX2DistilledTwoStagePipeline` | LTX-2 / LTX-2.3 merged-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + matching Lightricks upsampler | ✅︎ | ✅︎ | | | — |
|
||||
| `LTX2Pipeline` | LTX-2.5 Full/SFT one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
|
||||
| `LTX2TwoStagePipeline` | LTX-2.5 Full/SFT two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` + `Lightricks/LTX-2.5` LoRA | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
|
||||
| `LTX2DistilledOneStagePipeline` | LTX-2.5 merged-distilled one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
|
||||
| `LTX2DistilledTwoStagePipeline` | LTX-2.5 merged-distilled two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
|
||||
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
|
||||
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
|
||||
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) |
|
||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
|
||||
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
|
||||
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
@@ -58,6 +66,7 @@ th {
|
||||
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
|
||||
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
|
||||
| `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — |
|
||||
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
|
||||
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
|
||||
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
@@ -69,6 +78,7 @@ th {
|
||||
| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — |
|
||||
| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — |
|
||||
| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — |
|
||||
| `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
@@ -96,6 +106,7 @@ th {
|
||||
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | — |
|
||||
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
|
||||
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | | — |
|
||||
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
|
||||
|
||||
The `Recipe` column links the corresponding deployment recipe. `Published` is
|
||||
@@ -109,6 +120,10 @@ support metadata until their recipe is audited.
|
||||
|
||||
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
|
||||
|
||||
See the [LingBot-World 2.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Robbyant/LingBot-World-2.0.md)
|
||||
for offline and experimental realtime usage, supported controls, validation,
|
||||
and current limitations.
|
||||
|
||||
<sup>H3</sup> MiniMax H3 AMD GPU support is validated on gfx942 (MI300X) and
|
||||
gfx950 (MI350) in BF16, using the AITER `FLASH_ATTN` diffusion attention backend.
|
||||
Other AMD SKUs (e.g. MI325X) are not listed until their
|
||||
|
||||
@@ -141,7 +141,20 @@ The response shape depends on the streaming parameters:
|
||||
|
||||
**Non-streaming (default).** With `stream=false` and no `stream_format`, returns the
|
||||
complete clip as binary audio data with an appropriate `Content-Type` header (e.g.
|
||||
`audio/wav`). The raw-bytes body has no JSON carrier, so no `usage` is reported.
|
||||
`audio/wav`). Because the raw-bytes body has no JSON carrier, successful
|
||||
non-streaming responses from non-diffusion speech servers report usage through
|
||||
response headers:
|
||||
|
||||
| Header | Description |
|
||||
| --- | --- |
|
||||
| `x-vllm-omni-input-tokens` | Total input tokens (`text_tokens` + `audio_tokens`). |
|
||||
| `x-vllm-omni-output-tokens` | Generated codec/audio tokens. |
|
||||
| `x-vllm-omni-total-tokens` | `input_tokens` + `output_tokens`. |
|
||||
| `x-vllm-omni-input-text-tokens` | Tokens from the synthesized text (`input` plus `instructions`). |
|
||||
| `x-vllm-omni-input-audio-tokens` | Reference-audio codec frames, non-zero only for in-context voice cloning. |
|
||||
|
||||
Diffusion-mode speech servers route through a separate response path and do not
|
||||
emit these headers.
|
||||
|
||||
**Raw audio stream** (`stream_format="audio"`). Streams raw audio bytes (PCM or
|
||||
WAV) as they are decoded.
|
||||
|
||||
@@ -167,6 +167,14 @@ transfer options such as `control_path`, `control_guidance`,
|
||||
`num_conditional_frames`, `show_control_condition`, and `show_input`; see the
|
||||
Cosmos3 recipe for complete examples.
|
||||
|
||||
HTTP redirects for `image_reference.image_url` follow vLLM's
|
||||
`VLLM_MEDIA_URL_ALLOW_REDIRECTS` setting. Before starting the server, set it to
|
||||
`1` (the default) to allow redirects or `0` to reject them. A redirect target
|
||||
can differ from the original host, and vLLM-Omni does not yet fully implement
|
||||
upstream vLLM's media URL allowlist protection. Deployments should therefore
|
||||
treat remote media URLs as untrusted and choose this setting as part of their
|
||||
URL access policy.
|
||||
|
||||
### Speech-to-Video
|
||||
|
||||
For models that support audio-driven generation (e.g., Wan2.2-S2V), pass both
|
||||
|
||||
+3
-3
@@ -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.
|
||||
@@ -1,494 +1,192 @@
|
||||
# Diffusion Attention Backends
|
||||
|
||||
This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, the per-platform defaults, and how to use SageAttention.
|
||||
Use this page to select and configure a diffusion attention backend. Backend-
|
||||
specific installation, tuning, and compatibility details live in separate
|
||||
guides so that the selection contract stays easy to scan.
|
||||
|
||||
## Overview
|
||||
Diffusion attention backend selection applies to DiT attention in image and
|
||||
video generation models. It does **not** change autoregressive LLM attention,
|
||||
which uses vLLM's attention selector.
|
||||
|
||||
Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured.
|
||||
For the internal selector, registry, and platform contract, see
|
||||
[Attention Backend Selection](../../design/feature/attention_backend_selection.md).
|
||||
|
||||
This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. It does **not** affect autoregressive (LLM) attention paths — those go through vLLM's own attention backend selector.
|
||||
## Choose a guide
|
||||
|
||||
The full set of backends and their platform defaults is in the **Backend Options** and **Platform Defaults** sections below. If no attention backend is configured, vLLM-Omni asks the current platform to choose the default.
|
||||
| Need | Guide |
|
||||
| --- | --- |
|
||||
| Select a conservative or platform-native dense kernel | [Dense Backends](attention_backends/dense_backends.md) |
|
||||
| Use FlashInfer trtllm-gen FMHA, Skip-Softmax, or TRTLLM SAGE quantization | [TRTLLM Attention](attention_backends/trtllm.md) |
|
||||
| Install and use SageAttention 2.2 or SageAttention3 | [SageAttention](attention_backends/sage.md) |
|
||||
| Match training or rollout kernels loaded from Hugging Face | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
|
||||
| Use block-sparse video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
|
||||
|
||||
## Backend Options
|
||||
|
||||
| Value | Notes |
|
||||
|---|---|
|
||||
| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). Dense BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Packed paths can provide `cu_seqlens` directly. Supports optional **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
|
||||
| `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. |
|
||||
| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
|
||||
| `FLASHINFER_ATTN` | Uses FlashInfer's batch-prefill wrapper and supports mixed Q/K and V input dtypes through backend-specific configuration. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer` >= 0.6.16rc1 for mixed-dtype configurations. |
|
||||
| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
|
||||
| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
|
||||
| `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
|
||||
| `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. |
|
||||
| `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. |
|
||||
| `RAINFUSION_ATTN` | MindIE-SD **RainFusion** block-sparse video attention — see [below](#rainfusion_attn-backend-and-block-sparse-video-attention). Ascend NPU only; requires `mindiesd`. Delegates to `FLASH_ATTN` for anything that is not a packed video sequence. |
|
||||
## Backend options
|
||||
|
||||
| Value | Family | Primary use | Detail |
|
||||
| --- | --- | --- | --- |
|
||||
| `TORCH_SDPA` | Dense | Conservative reference; always available | [Dense Backends](attention_backends/dense_backends.md#torch_sdpa) |
|
||||
| `FLASH_ATTN` | Dense | FlashAttention 4/3/2 depending on the installed package and GPU | [Dense Backends](attention_backends/dense_backends.md#flash_attn) |
|
||||
| `CUDNN_ATTN` | Dense | Mask-heavy DiTs on Blackwell with cuDNN 9.5 or newer | [Dense Backends](attention_backends/dense_backends.md#cudnn_attn) |
|
||||
| `FLASHINFER_ATTN` | Dense or quantized | FlashInfer batch prefill; optional mixed Q/K and V dtypes | [Dense Backends](attention_backends/dense_backends.md#flashinfer_attn) |
|
||||
| `TRTLLM_ATTN` | Dense, sparse, or quantized | Datacenter Blackwell with `head_dim=128` and compatible packed paths | [TRTLLM Attention](attention_backends/trtllm.md) |
|
||||
| `SAGE_ATTN` | Quantized | SageAttention 2.2 INT8 attention | [SageAttention](attention_backends/sage.md#sage_attn) |
|
||||
| `SAGE_ATTN_3` | Quantized | SageAttention3 on Blackwell | [SageAttention](attention_backends/sage.md#sage_attn_3) |
|
||||
| `FLASH_ATTN_HUB` | Hub kernel | FlashAttention 2 from Hugging Face `kernels` | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
|
||||
| `FLASH_ATTN_3_HUB` | Hub kernel | FlashAttention 3 from Hugging Face `kernels` on Hopper or newer | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
|
||||
| `RAINFUSION_ATTN` | Block sparse | MindIE-SD RainFusion video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
|
||||
|
||||
## Configuration
|
||||
|
||||
Diffusion attention backends can be configured three ways, in priority order:
|
||||
Backend selection is resolved in this order:
|
||||
|
||||
1. **`--diffusion-attention-config`** — structured per-role config (highest priority).
|
||||
2. **`--diffusion-attention-backend` / `DIFFUSION_ATTENTION_BACKEND` env var** — global shorthand that sets the default backend.
|
||||
3. **Platform default** — used when nothing is configured.
|
||||
1. `--diffusion-attention-config` per-role configuration.
|
||||
2. `--diffusion-attention-backend` or `DIFFUSION_ATTENTION_BACKEND` as a
|
||||
global default.
|
||||
3. The current platform's default.
|
||||
|
||||
`--diffusion-attention-backend` is shorthand for `--diffusion-attention-config.default.backend`. It may be combined with `--diffusion-attention-config.per_role.*` overrides, but is mutually exclusive with `--diffusion-attention-config.default.backend`.
|
||||
`--diffusion-attention-backend` is shorthand for
|
||||
`--diffusion-attention-config.default.backend`. Do not pass it together with an
|
||||
explicit `default.backend` in the structured configuration.
|
||||
|
||||
### Global default
|
||||
|
||||
Set the default backend for every diffusion attention layer:
|
||||
|
||||
```bash
|
||||
# CLI flag
|
||||
vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN
|
||||
|
||||
# Environment variable (also recognized for backwards compatibility)
|
||||
export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN
|
||||
# Backwards-compatible environment variable
|
||||
export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
```
|
||||
|
||||
### Per-role configuration
|
||||
|
||||
Roles are free-form strings declared by each diffusion model. The two common categories are `"self"` and `"cross"`; model-specific roles (e.g. `"ltx2.audio_to_video"`) may also be declared. A role string is matched in this order:
|
||||
Roles are declared by each diffusion model. Common categories are `self` and
|
||||
`cross`; a model may also use a more specific role such as
|
||||
`ltx2.audio_to_video`. Resolution order is:
|
||||
|
||||
1. Exact `per_role[role]` match
|
||||
2. `per_role[role_category]` fallback (e.g. `"ltx2.audio_to_video"` → `"cross"`)
|
||||
3. `default`
|
||||
4. Platform default
|
||||
|
||||
Use vLLM-style dotted flags or one JSON blob:
|
||||
1. Exact `per_role[role]` match.
|
||||
2. Category `per_role[role_category]` match.
|
||||
3. `default`.
|
||||
4. Platform default.
|
||||
|
||||
```bash
|
||||
# Dotted flags
|
||||
vllm-omni serve <model> \
|
||||
--diffusion-attention-config.default.backend FLASH_ATTN \
|
||||
--diffusion-attention-config.per_role.cross.backend TORCH_SDPA
|
||||
--diffusion-attention-config.default.backend FLASH_ATTN \
|
||||
--diffusion-attention-config.per_role.cross.backend TORCH_SDPA
|
||||
|
||||
# JSON
|
||||
# Equivalent JSON
|
||||
vllm-omni serve <model> \
|
||||
--diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
|
||||
--diffusion-attention-config \
|
||||
'{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
|
||||
```
|
||||
|
||||
A backend that needs configuration exposes it as a typed field on the spec:
|
||||
`TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)) and
|
||||
`RAINFUSION_ATTN`'s RainFusion block (see [below](#rainfusion_attn-backend-and-block-sparse-video-attention)):
|
||||
|
||||
```bash
|
||||
--diffusion-attention-config.default.backend TRTLLM_ATTN \
|
||||
--diffusion-attention-config.default.skip_softmax.target_sparsity 0.5
|
||||
```
|
||||
|
||||
### Programmatic API
|
||||
|
||||
When constructing `OmniDiffusionConfig` directly:
|
||||
### Python API
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.data import (
|
||||
AttentionConfig,
|
||||
AttentionSpec,
|
||||
OmniDiffusionConfig,
|
||||
)
|
||||
|
||||
config = OmniDiffusionConfig(
|
||||
diffusion_attention_config=AttentionConfig(
|
||||
default=AttentionSpec(backend="FLASH_ATTN"),
|
||||
per_role={
|
||||
"cross": AttentionSpec(backend="TORCH_SDPA"),
|
||||
},
|
||||
per_role={"cross": AttentionSpec(backend="TORCH_SDPA")},
|
||||
),
|
||||
...,
|
||||
)
|
||||
```
|
||||
|
||||
A plain dict is also accepted and normalized to `AttentionConfig`.
|
||||
Backend-specific typed blocks are documented with their consumers:
|
||||
|
||||
#### FlashInfer QK16/V8 quantized attention
|
||||
- `quant`: [FlashInfer](attention_backends/dense_backends.md#flashinfer-quantized-attention)
|
||||
and [TRTLLM SAGE](attention_backends/trtllm.md#sage-quantization).
|
||||
- `skip_softmax`: [TRTLLM Skip-Softmax](attention_backends/trtllm.md#skip-softmax).
|
||||
- `block_sparse`: [RainFusion](attention_backends/rainfusion.md#configuration).
|
||||
|
||||
`FLASHINFER_ATTN` accepts quantization options through `AttentionSpec.quant`. For QK16/V8:
|
||||
|
||||
```python
|
||||
config = OmniDiffusionConfig(
|
||||
diffusion_attention_config=AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="FLASHINFER_ATTN",
|
||||
quant=AttnQuantSpec(
|
||||
dtype_qk="bfloat16",
|
||||
dtype_vo="fp8_e4m3",
|
||||
),
|
||||
),
|
||||
),
|
||||
...,
|
||||
)
|
||||
```
|
||||
|
||||
`dtype_qk` configures the Q and K input dtypes. `dtype_vo` configures the V input dtype.
|
||||
|
||||
## Platform Defaults
|
||||
## Platform defaults
|
||||
|
||||
### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
|
||||
|
||||
Auto-route preference, in order:
|
||||
The CUDA auto-route preference is:
|
||||
|
||||
1. `TRTLLM_ATTN` — on **datacenter** Blackwell (sm_100 / sm_103) when `flashinfer` is installed, the model's `head_dim` is 128, **and the model declares a compatible packed/mask-free path**
|
||||
2. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
|
||||
3. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
|
||||
4. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
|
||||
5. `TORCH_SDPA` — last resort
|
||||
1. `TRTLLM_ATTN` on datacenter Blackwell (sm_100/sm_103) when FlashInfer is
|
||||
available, `head_dim=128`, and the model declares a compatible packed or
|
||||
mask-free path.
|
||||
2. `CUDNN_ATTN` when cuDNN 9.5 or newer is available.
|
||||
3. `FLASHINFER_ATTN` when FlashInfer is available but cuDNN is too old.
|
||||
4. `FLASH_ATTN` when a compatible package is installed.
|
||||
5. `TORCH_SDPA`.
|
||||
|
||||
`TRTLLM_ATTN` is skipped on workstation Blackwell (sm_120 / sm_121) and for any `head_dim != 128`, so those GPUs keep the `CUDNN_ATTN` route described below.
|
||||
`TRTLLM_ATTN` is not auto-selected on workstation Blackwell
|
||||
(sm_120/sm_121), for other head dimensions, or for paths that require an
|
||||
unsupported mask.
|
||||
|
||||
`TRTLLM_ATTN` outranks `CUDNN_ATTN` on datacenter Blackwell for compatible packed/mask-free pipelines. Workstation Blackwell (sm_120 / sm_121) and pipelines that require attention masks retain their normal fallback.
|
||||
### Hopper, Ada, and Ampere
|
||||
|
||||
The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.
|
||||
The CUDA auto-route uses `FLASH_ATTN` when available and otherwise falls back
|
||||
to `TORCH_SDPA`. `CUDNN_ATTN` and `FLASHINFER_ATTN` remain explicit options.
|
||||
|
||||
**Why CUDNN_ATTN by default**: on mask-heavy diffusion models (HunyuanVideo-1.5, Qwen-Image), cuDNN's pinned FMHA kernel sidesteps a PyTorch SDPA dispatch quirk where the unpinned dispatcher picks `EFFICIENT_ATTENTION` (~25 ms) for masked calls instead of cuDNN (~11 ms). The pin gives 2× e2e on HV-1.5 with no regression on lighter models.
|
||||
Other platforms validate an explicit backend and choose their own default
|
||||
through the platform implementation. Check the startup log to confirm the
|
||||
resolved backend.
|
||||
|
||||
### Hopper (sm_90) / Ada (sm_89) / Ampere (sm_80–sm_86)
|
||||
## Choosing a backend manually
|
||||
|
||||
Auto-route preference:
|
||||
Override the platform default when you need:
|
||||
|
||||
1. `FLASH_ATTN` — when `flash-attn` is installed
|
||||
2. `TORCH_SDPA` — fallback
|
||||
- a correctness reference (`TORCH_SDPA`);
|
||||
- a backend-specific workaround;
|
||||
- training/rollout kernel alignment (Hub backends); or
|
||||
- an explicitly validated sparse or quantized speedup.
|
||||
|
||||
`CUDNN_ATTN` and `FLASHINFER_ATTN` are still selectable via env var on these GPUs but are not in the auto-route — FlashAttention 2 is the well-tuned path on pre-Blackwell hardware.
|
||||
The startup log prints the resolved backend and whether it came from explicit
|
||||
configuration or platform defaulting. If no resolution message appears, check
|
||||
earlier logs for diffusion-stage initialization failures.
|
||||
|
||||
## TRTLLM_ATTN Backend and Skip-Softmax
|
||||
## Reference benchmark
|
||||
|
||||
`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA. Dense it is a BF16 backend like the others; on top of
|
||||
that it can enable **Skip-Softmax**, a sparse-attention mode that trades a little fidelity for speed
|
||||
(algorithm: [design doc](../../design/feature/skip_softmax.md)).
|
||||
The following BF16 results were measured on an sm_120 RTX Pro 6000 Blackwell
|
||||
with the same prompt and seed across runs. Treat them as reference results, not
|
||||
portable guarantees.
|
||||
|
||||
Enable it through the typed `skip_softmax` block on the attention spec:
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
|---|---|---|
|
||||
| `target_sparsity` | finite, `[0, 1]` | Operating point on a calibrated curve. Needs a calibration for the model. |
|
||||
| `threshold` | finite, `≥ 0` | Direct threshold, no calibration needed. Mutually exclusive with `target_sparsity`. |
|
||||
| `disabled_until_timestep` | finite, `[0, 1]` | Holds early, high-noise steps dense; skip turns on once the normalized timestep `t ≤ D`. |
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
|
||||
"skip_softmax": {"target_sparsity": 0.65, "disabled_until_timestep": 0.86}}}'
|
||||
```
|
||||
|
||||
Programmatically the same block is a typed `SkipSoftmaxSpec` (values validated at construction):
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
|
||||
|
||||
AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
skip_softmax=SkipSoftmaxSpec(target_sparsity=0.65, disabled_until_timestep=0.86),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
**Start at `target_sparsity=0.65, disabled_until_timestep=0.86`** and raise `target_sparsity` only
|
||||
as far as your quality bar allows — the output diverges from dense as it climbs. The gain grows with
|
||||
sequence length, since Skip-Softmax only accelerates attention. Dense BF16 is `TRTLLM_ATTN` with no
|
||||
`skip_softmax` block.
|
||||
|
||||
Requires datacenter Blackwell with `head_dim == 128`; elsewhere, or without FlashInfer, selecting
|
||||
`TRTLLM_ATTN` raises rather than silently degrading.
|
||||
|
||||
## TRTLLM_ATTN SAGE Quantization
|
||||
|
||||
On top of dense/Skip-Softmax, `TRTLLM_ATTN` can run **SAGE** (SageAttention) — the trtllm-gen
|
||||
FMHA kernel with per-block quantized Q/K and per-channel fp8 V, so the two attention matmuls run
|
||||
in low precision. Enable it through the typed `quant` block on the attention spec:
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
|---|---|---|
|
||||
| `dtype_qk` | `int8`, `fp8_e4m3` | Quantization dtype for Q and K. Absent ⇒ dense (SAGE off). |
|
||||
| `q_block_size` | `1`, `4`, `16` | Per-token block size for Q scales (default `1`). |
|
||||
| `k_block_size` | `1`, `4`, `16` | Per-token block size for K scales (default `16`). |
|
||||
|
||||
V is always quantized per-channel to fp8_e4m3 and K-smoothing is applied inside the routine, so
|
||||
neither is a user knob. Only block sizes `1`, `4`, `16` have compiled kernels.
|
||||
|
||||
The `quant` block is a shared `AttnQuantSpec` consumed by more than one backend; each reads the
|
||||
fields that apply and rejects the rest. `TRTLLM_ATTN` (SAGE) reads `dtype_qk` (`int8` / `fp8_e4m3`),
|
||||
`q_block_size`, and `k_block_size`; `FLASHINFER_ATTN` reads `dtype_qk` (`float16` / `bfloat16`),
|
||||
`dtype_vo`, and `flashinfer_backend` (see its section). A `dtype_qk` from the wrong set raises.
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
|
||||
"quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
|
||||
|
||||
AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
quant=AttnQuantSpec(dtype_qk="fp8_e4m3", q_block_size=1, k_block_size=16),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
**Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel
|
||||
availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103**
|
||||
(B300); `int8` QK kernels are compiled for **SM100 only**.
|
||||
|
||||
## RAINFUSION_ATTN Backend and Block-Sparse Video Attention
|
||||
|
||||
`RAINFUSION_ATTN` runs MindIE-SD's RainFusion (`rf_v2`) kernel on Ascend NPU. It pools each
|
||||
128-token block of keys, ranks them per query block, and computes attention against only the
|
||||
top-scoring ones. Video tokens are rearranged into `(t, h, w)` order first, so the kept blocks are
|
||||
spatiotemporal neighbours rather than arbitrary rows of the packed sequence.
|
||||
|
||||
Only the video segment is sparsified. The prefix rows (text, visual conditions, audio) and the
|
||||
first-frame blocks are always kept dense, which is why the realized sparsity is lower than the
|
||||
nominal `sparsity` you configure. Anything the kernel cannot handle — a warmup denoise step, an
|
||||
exempt layer, a sequence with no published video geometry, a video segment under 32 blocks —
|
||||
delegates to `FLASH_ATTN`, so a model can select this backend unconditionally.
|
||||
|
||||
Enable it through the typed `block_sparse` block on the attention spec, shared by every
|
||||
block-sparse backend:
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
|---|---|---|
|
||||
| `sparsity` | finite, `[0, 1]` | Nominal fraction of key blocks dropped per query block. `0` disables sparsity. Defaults to `0.8`. |
|
||||
| `start_step` | `≥ 0` | Keep the first N denoise steps dense. Layout is decided early, so these steps dominate structural fidelity. |
|
||||
| `skip_layers` | index selector, e.g. `"0-3,38"` | DiT blocks that always stay dense. |
|
||||
|
||||
```bash
|
||||
vllm-omni serve MiniMaxAI/MiniMax-H3 \
|
||||
--diffusion-attention-config '{"default": {"backend": "RAINFUSION_ATTN",
|
||||
"block_sparse": {"sparsity": 0.8, "start_step": 0}}}'
|
||||
```
|
||||
|
||||
Programmatically the same block is a typed `BlockSparseSpec` (values validated at construction):
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec
|
||||
|
||||
AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="RAINFUSION_ATTN",
|
||||
block_sparse=BlockSparseSpec(sparsity=0.8, start_step=0, skip_layers="0-1"),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
**Tune in the order `start_step` → `sparsity` → `skip_layers`.** Raise `start_step` first: it is the
|
||||
cheapest way to recover structure, because the early high-noise steps decide global layout while the
|
||||
later steps only refine texture. Then walk `sparsity` up to your quality bar. Reach for
|
||||
`skip_layers` last, once an A/B against dense at the same seed points at specific blocks.
|
||||
|
||||
Requires Ascend NPU with `mindiesd`; selecting it on any other platform raises. It is also
|
||||
incompatible with ring sequence parallelism, since `rf_v2` needs the whole key sequence to rank
|
||||
blocks — use Ulysses SP (`ring_degree=1`).
|
||||
|
||||
### Which geometries run sparse
|
||||
|
||||
MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when
|
||||
necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2
|
||||
does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the
|
||||
promoted rows remain visible to every query and key.
|
||||
|
||||
For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row
|
||||
128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix.
|
||||
MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask
|
||||
and kernel tiling aligned, at the cost of retaining more blocks.
|
||||
|
||||
8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8
|
||||
creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent
|
||||
grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release
|
||||
that implements protected video tails for rf_v2.
|
||||
|
||||
## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
|
||||
|
||||
Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
|
||||
|
||||
| Model | Shape | TORCH_SDPA | CUDNN_ATTN | FLASHINFER_ATTN |
|
||||
|---|---|---|---|---|
|
||||
| Model | Shape | `TORCH_SDPA` | `CUDNN_ATTN` | `FLASHINFER_ATTN` |
|
||||
| --- | --- | ---: | ---: | ---: |
|
||||
| HunyuanVideo-1.5 (T2V) | 480p / 33f / 50 steps | 147.05 s | **73.02 s** | 127.84 s |
|
||||
| Wan 2.2 14B (T2V) | 480p / 33f / 40 steps | 117.75 s | 117.17 s | **115.07 s** |
|
||||
| Qwen-Image (T2I) | 1024² / 50 steps | 17.41 s | **15.14 s** | 16.02 s |
|
||||
| FLUX.2-dev (T2I) | 1024² / 50 steps, TP=2 | 53.62 s | **53.30 s** | 54.94 s |
|
||||
|
||||
Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
|
||||
Mask-heavy DiTs favored `CUDNN_ATTN`; lighter-mask or TP-saturated workloads
|
||||
were close enough that users should benchmark their exact model and shape.
|
||||
|
||||
## Known Limitations
|
||||
## Compatibility anchors
|
||||
|
||||
### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
|
||||
The following headings preserve links to sections that moved into dedicated
|
||||
guides.
|
||||
|
||||
LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
|
||||
## TRTLLM_ATTN Backend and Skip-Softmax
|
||||
|
||||
**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
|
||||
See [TRTLLM Attention](attention_backends/trtllm.md#skip-softmax).
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Lightricks/LTX-2 ...
|
||||
```
|
||||
## TRTLLM_ATTN SAGE Quantization
|
||||
|
||||
### FlashAttention-4 on Blackwell
|
||||
See [TRTLLM Attention](attention_backends/trtllm.md#sage-quantization).
|
||||
|
||||
Install the optional CUDA 13 extra to use the CuTe-based FA4 path:
|
||||
## RAINFUSION_ATTN Backend and Block-Sparse Video Attention
|
||||
|
||||
```bash
|
||||
pip install 'vllm-omni[fa4]'
|
||||
```
|
||||
|
||||
`FLASH_ATTN` prefers `flash_attn.cute` on Blackwell and falls back to FA3/FA2
|
||||
when it is unavailable. Version `4.0.0b18` is required; earlier beta
|
||||
wheels had known JIT failures on Blackwell.
|
||||
|
||||
## Choosing a Backend Manually
|
||||
|
||||
### When to override the default
|
||||
|
||||
- **Quality validation**: compare a new backend against `TORCH_SDPA` as the reference, since SDPA's default dispatcher is the most extensively tested.
|
||||
- **Lossy speedup hunting**: try `SAGE_ATTN` (INT8 quantized) on diffusion outputs — typically indistinguishable visually but always validate.
|
||||
- **Workaround for known issues**: see Known Limitations above.
|
||||
|
||||
### Verifying which backend is in use
|
||||
|
||||
The startup log prints one of:
|
||||
|
||||
```
|
||||
Using diffusion attention backend 'CUDNN_ATTN' # explicit override
|
||||
Defaulting to diffusion attention backend CUDNN_ATTN ... # auto-route
|
||||
Defaulting to diffusion attention backend SDPA # nothing else available
|
||||
```
|
||||
|
||||
If you don't see one of these, the model didn't reach diffusion stage init — check earlier logs for failures.
|
||||
See [RainFusion](attention_backends/rainfusion.md).
|
||||
|
||||
## SageAttention Installation
|
||||
|
||||
vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
|
||||
|
||||
Build from source:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/thu-ml/SageAttention.git
|
||||
cd SageAttention
|
||||
|
||||
export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
|
||||
pip install . --no-build-isolation
|
||||
```
|
||||
|
||||
Quick check:
|
||||
|
||||
```bash
|
||||
python -c "import sageattention; print(sageattention.__file__)"
|
||||
```
|
||||
See [SageAttention](attention_backends/sage.md#installation).
|
||||
|
||||
## SageAttention3 Installation
|
||||
|
||||
vLLM-Omni expects SageAttention3 to be installed into the same Python environment as vLLM-Omni.
|
||||
|
||||
Build from source:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/thu-ml/SageAttention.git
|
||||
cd SageAttention/sageattention3_blackwell
|
||||
python setup.py install
|
||||
```
|
||||
|
||||
Quick check:
|
||||
|
||||
```bash
|
||||
python -c "import sageattn3; print(sageattn3.__file__)"
|
||||
```
|
||||
|
||||
Notes:
|
||||
|
||||
- `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class.
|
||||
- SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness.
|
||||
See [SageAttention](attention_backends/sage.md#sageattention3-installation).
|
||||
|
||||
## HuggingFace Kernels Hub Backends
|
||||
|
||||
To achieve perfect numerical consistency between **training** (typically using HuggingFace Diffusers with Hub kernels) and **serving/rollout** (using vLLM-Omni), you can use the Hub-based attention backends. This eliminates numerical drift / sampling divergence caused by executing different local kernel versions during rollout.
|
||||
|
||||
The following backend options are supported:
|
||||
- `FLASH_ATTN_HUB` (HuggingFace `kernels-community/flash-attn2`)
|
||||
- `FLASH_ATTN_3_HUB` (HuggingFace `kernels-community/flash-attn3`, Hopper sm_90+ only)
|
||||
|
||||
### Installation
|
||||
|
||||
To use these backends, you must install the `kernels` library:
|
||||
|
||||
```bash
|
||||
pip install kernels==0.14.1
|
||||
```
|
||||
|
||||
If the `kernels` library is not available in the environment, vLLM-Omni will log a warning and fall back gracefully to the corresponding local backend implementations (`FLASH_ATTN`). On CUDA GPUs below Hopper (compute capability < 9.0), `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
|
||||
|
||||
### Usage
|
||||
|
||||
Select a Hub backend using the global CLI flag or environment variables:
|
||||
|
||||
```bash
|
||||
# Environment variable
|
||||
export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB
|
||||
|
||||
# CLI flag
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN_HUB
|
||||
```
|
||||
|
||||
## Usage Examples
|
||||
|
||||
### Default (auto-route)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 50 --seed 42 --guidance-scale 6.0 \
|
||||
--output hv15.mp4
|
||||
```
|
||||
|
||||
On Blackwell this picks `CUDNN_ATTN` automatically. Check the log for the `Defaulting to ...` line.
|
||||
|
||||
### Explicit backend selection
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Lightricks/LTX-2 \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 40 --seed 42 --guidance-scale 4.0 \
|
||||
--output ltx2.mp4
|
||||
```
|
||||
|
||||
### SageAttention (lossy)
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output hv15_sage.mp4
|
||||
```
|
||||
|
||||
Example: Wan2.2 TI2V 5B
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 704 --width 1280 --num-frames 49 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 5.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output outputs/wan22_sage.mp4
|
||||
```
|
||||
|
||||
### Enable SageAttention3
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output outputs/hv15_sage3.mp4
|
||||
```
|
||||
See [Hugging Face Hub Backends](attention_backends/huggingface_hub.md).
|
||||
@@ -0,0 +1,110 @@
|
||||
# Dense Attention Backends
|
||||
|
||||
Use these backends when you want dense attention without backend-specific
|
||||
sparsity. Start with `TORCH_SDPA` for correctness comparisons, then benchmark
|
||||
the fastest compatible kernel for your model, shape, and hardware.
|
||||
|
||||
For selection precedence and per-role configuration, see the
|
||||
[attention backend overview](../attention_backends.md).
|
||||
|
||||
## `TORCH_SDPA`
|
||||
|
||||
`TORCH_SDPA` calls PyTorch `scaled_dot_product_attention` and lets PyTorch's
|
||||
dispatcher choose the implementation. It is always available and is the most
|
||||
conservative reference when validating another backend.
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend TORCH_SDPA
|
||||
```
|
||||
|
||||
## `FLASH_ATTN`
|
||||
|
||||
`FLASH_ATTN` uses the installed FlashAttention implementation. On Blackwell it
|
||||
prefers the CuTe-based FlashAttention 4 path and falls back to
|
||||
FlashAttention 3 or 2 when that path is unavailable. On Hopper, Ada, and
|
||||
Ampere it is the preferred automatic route when a compatible package is
|
||||
installed.
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN
|
||||
```
|
||||
|
||||
### FlashAttention 4 on Blackwell
|
||||
|
||||
Install the optional CUDA 13 extra:
|
||||
|
||||
```bash
|
||||
pip install 'vllm-omni[fa4]'
|
||||
```
|
||||
|
||||
Version `4.0.0b18` is required; earlier beta wheels had known JIT failures on
|
||||
Blackwell. If the CuTe path is unavailable, the backend falls back to the
|
||||
compatible FlashAttention 3 or 2 path.
|
||||
|
||||
## `CUDNN_ATTN`
|
||||
|
||||
`CUDNN_ATTN` pins PyTorch SDPA to `CUDNN_ATTENTION`. It is particularly useful
|
||||
for mask-heavy DiTs and is automatically preferred on Blackwell when cuDNN
|
||||
9.5 or newer is available and the higher-priority TRTLLM route is not
|
||||
compatible.
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend CUDNN_ATTN
|
||||
```
|
||||
|
||||
### LTX-2.0 limitation
|
||||
|
||||
LTX-2 audio attention has a symbolic head dimension during `torch.compile`
|
||||
tracing. The cuDNN SDPA selector rejects that symbolic dimension and Dynamo
|
||||
aborts compilation. This is tracked in
|
||||
[issue #3121](https://github.com/vllm-project/vllm-omni/issues/3121).
|
||||
|
||||
Use `FLASHINFER_ATTN` or `TORCH_SDPA` as a workaround:
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN \
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Lightricks/LTX-2 ...
|
||||
```
|
||||
|
||||
## `FLASHINFER_ATTN`
|
||||
|
||||
`FLASHINFER_ATTN` uses FlashInfer's batch-prefill wrapper. It is an explicit
|
||||
option on CUDA platforms and an automatic Blackwell fallback when FlashInfer
|
||||
is installed but cuDNN is too old for `CUDNN_ATTN`.
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASHINFER_ATTN
|
||||
```
|
||||
|
||||
### FlashInfer quantized attention
|
||||
|
||||
The backend accepts an `AttentionSpec.quant` block. For QK16/V8, keep Q and K
|
||||
in FP16 or BF16 and use FP8 E4M3 for V:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import (
|
||||
AttentionConfig,
|
||||
AttentionSpec,
|
||||
AttnQuantSpec,
|
||||
OmniDiffusionConfig,
|
||||
)
|
||||
|
||||
config = OmniDiffusionConfig(
|
||||
diffusion_attention_config=AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="FLASHINFER_ATTN",
|
||||
quant=AttnQuantSpec(
|
||||
dtype_qk="bfloat16",
|
||||
dtype_vo="fp8_e4m3",
|
||||
),
|
||||
),
|
||||
),
|
||||
...,
|
||||
)
|
||||
```
|
||||
|
||||
`dtype_qk` controls Q and K; `dtype_vo` controls V. Mixed-dtype
|
||||
configurations require FlashInfer 0.6.16rc1 or newer. The shared quantization
|
||||
schema is also consumed by TRTLLM, but each backend validates its own allowed
|
||||
fields and values; see [TRTLLM SAGE quantization](trtllm.md#sage-quantization).
|
||||
@@ -0,0 +1,32 @@
|
||||
# Hugging Face Hub Attention Backends
|
||||
|
||||
Hub backends load the same published kernels commonly used with Hugging Face
|
||||
Diffusers. They are useful when training and serving must use matching kernel
|
||||
implementations to reduce numerical drift and sampling divergence.
|
||||
|
||||
| Backend | Kernel | Platform |
|
||||
| --- | --- | --- |
|
||||
| `FLASH_ATTN_HUB` | `kernels-community/flash-attn2` | Compatible CUDA GPUs |
|
||||
| `FLASH_ATTN_3_HUB` | `kernels-community/flash-attn3` | Hopper sm_90 or newer |
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install kernels==0.14.1
|
||||
```
|
||||
|
||||
If `kernels` is unavailable, vLLM-Omni warns and falls back to the
|
||||
corresponding local FlashAttention implementation. On CUDA devices older than
|
||||
Hopper, `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
|
||||
|
||||
## Usage
|
||||
|
||||
```bash
|
||||
export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB
|
||||
|
||||
# Equivalent CLI selection
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN_HUB
|
||||
```
|
||||
|
||||
For per-role kernel alignment, configure the Hub backend on only the relevant
|
||||
role; see the [attention backend overview](../attention_backends.md#per-role-configuration).
|
||||
@@ -0,0 +1,72 @@
|
||||
# RainFusion Attention
|
||||
|
||||
`RAINFUSION_ATTN` runs MindIE-SD RainFusion (`rf_v2`) block-sparse video
|
||||
attention on Ascend NPU. It pools 128-token key blocks, ranks them per query
|
||||
block, and attends to the highest-scoring blocks after arranging video tokens
|
||||
in `(t, h, w)` order.
|
||||
|
||||
Only the video segment is sparse. Prefix rows such as text, visual conditions,
|
||||
and audio, plus first-frame blocks, remain dense. Unsupported calls—including
|
||||
warmup steps, skipped layers, missing video geometry, and video segments below
|
||||
32 blocks—delegate to `FLASH_ATTN`, so compatible models can select RainFusion
|
||||
globally.
|
||||
|
||||
## Configuration
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
| --- | --- | --- |
|
||||
| `sparsity` | finite, `[0, 1]` | Nominal dropped-key-block fraction; default `0.8`; `0` disables sparsity |
|
||||
| `start_step` | integer, `>= 0` | Number of early denoise steps kept dense |
|
||||
| `skip_layers` | selector such as `"0-3,38"` | DiT blocks kept dense |
|
||||
|
||||
```bash
|
||||
vllm-omni serve MiniMaxAI/MiniMax-H3 \
|
||||
--diffusion-attention-config '{"default":{"backend":"RAINFUSION_ATTN",\
|
||||
"block_sparse":{"sparsity":0.8,"start_step":0}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec
|
||||
|
||||
config = AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="RAINFUSION_ATTN",
|
||||
block_sparse=BlockSparseSpec(
|
||||
sparsity=0.8,
|
||||
start_step=0,
|
||||
skip_layers="0-1",
|
||||
),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
Tune `start_step`, then `sparsity`, then `skip_layers`. Increasing the dense
|
||||
early-step window is usually the cheapest way to recover global structure;
|
||||
use layer exclusions only after a same-seed dense comparison identifies
|
||||
sensitive blocks.
|
||||
|
||||
## Requirements and compatibility
|
||||
|
||||
RainFusion requires Ascend NPU and `mindiesd`. Selecting it on another
|
||||
platform raises. It is incompatible with ring sequence parallelism because
|
||||
the kernel needs the complete key sequence for ranking; use Ulysses sequence
|
||||
parallelism with `ring_degree=1`.
|
||||
|
||||
## Geometry handling
|
||||
|
||||
RainFusion handles arbitrary video grids by rearranging video spatially and,
|
||||
when necessary, promoting a real-video suffix to the always-kept prefix. It
|
||||
does not pad because `rf_v2` does not consume a padding attention mask.
|
||||
|
||||
For MiniMax-H3 at 1344x768, grid `(62, 24, 42)` produces 62,496 video rows.
|
||||
The implementation promotes 2,976 rows and leaves 59,520 sparse rows
|
||||
(`465 x 128`) so the mask and kernel tiling remain aligned.
|
||||
|
||||
Spatial grids aligned to 8x8 are still preferable. A latent height or width
|
||||
not divisible by 8 creates an always-kept suffix and reduces realized
|
||||
sparsity. At the image level, that means width and height divisible by 256.
|
||||
Protected video tails require a compatible MindIE-SD release.
|
||||
|
||||
For common configuration and selector behavior, see the
|
||||
[attention backend overview](../attention_backends.md) and the
|
||||
[backend selection design](../../../design/feature/attention_backend_selection.md).
|
||||
@@ -0,0 +1,63 @@
|
||||
# SageAttention
|
||||
|
||||
SageAttention backends provide lossy low-precision attention for diffusion
|
||||
models. Validate output quality against `TORCH_SDPA` at the same seed before
|
||||
using either backend in production.
|
||||
|
||||
## `SAGE_ATTN`
|
||||
|
||||
`SAGE_ATTN` uses SageAttention 2.2 with INT8-quantized attention and FP16
|
||||
accumulation.
|
||||
|
||||
### Installation
|
||||
|
||||
Install SageAttention into the same environment as vLLM-Omni:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/thu-ml/SageAttention.git
|
||||
cd SageAttention
|
||||
export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
|
||||
pip install . --no-build-isolation
|
||||
```
|
||||
|
||||
Verify the installation:
|
||||
|
||||
```bash
|
||||
python -c "import sageattention; print(sageattention.__file__)"
|
||||
```
|
||||
|
||||
Select it globally:
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN
|
||||
```
|
||||
|
||||
## `SAGE_ATTN_3`
|
||||
|
||||
`SAGE_ATTN_3` uses the SageAttention3 Blackwell implementation.
|
||||
|
||||
### SageAttention3 installation
|
||||
|
||||
```bash
|
||||
git clone https://github.com/thu-ml/SageAttention.git
|
||||
cd SageAttention/sageattention3_blackwell
|
||||
python setup.py install
|
||||
```
|
||||
|
||||
Verify the installation:
|
||||
|
||||
```bash
|
||||
python -c "import sageattn3; print(sageattn3.__file__)"
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN_3
|
||||
```
|
||||
|
||||
`SAGE_ATTN_3` requires CUDA, an importable `sageattn3`, and a Blackwell-class
|
||||
GPU. Its kernel assumes the query-head count equals the key/value-head count.
|
||||
GQA and MQA diffusion calls therefore fall back to PyTorch SDPA for
|
||||
correctness.
|
||||
|
||||
For common configuration and platform routing, see the
|
||||
[attention backend overview](../attention_backends.md).
|
||||
@@ -0,0 +1,103 @@
|
||||
# TRTLLM Attention
|
||||
|
||||
`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels. It supports dense
|
||||
BF16 attention plus optional Skip-Softmax sparsity or SAGE quantization.
|
||||
|
||||
For common selection and per-role configuration, see the
|
||||
[attention backend overview](../attention_backends.md).
|
||||
|
||||
## Requirements
|
||||
|
||||
Selecting `TRTLLM_ATTN` requires:
|
||||
|
||||
- datacenter Blackwell, sm_100 or sm_103;
|
||||
- `head_dim=128`;
|
||||
- FlashInfer; and
|
||||
- a model-declared packed or mask-free compatible path.
|
||||
|
||||
It is not supported on workstation Blackwell (sm_120/sm_121). An incompatible
|
||||
explicit selection raises instead of silently changing the backend.
|
||||
|
||||
## Dense mode
|
||||
|
||||
Select `TRTLLM_ATTN` without `skip_softmax` or `quant` to run dense BF16:
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend TRTLLM_ATTN
|
||||
```
|
||||
|
||||
## Skip-Softmax
|
||||
|
||||
Skip-Softmax trades some fidelity for speed by skipping low-contribution
|
||||
attention work. See the [feature design](../../../design/feature/skip_softmax.md)
|
||||
for the algorithm and implementation contract.
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
| --- | --- | --- |
|
||||
| `target_sparsity` | finite, `[0, 1]` | Point on a model-calibrated curve |
|
||||
| `threshold` | finite, `>= 0` | Direct threshold; mutually exclusive with `target_sparsity` |
|
||||
| `disabled_until_timestep` | finite, `[0, 1]` | Keep early high-noise steps dense until normalized `t <= D` |
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
|
||||
"skip_softmax":{"target_sparsity":0.65,"disabled_until_timestep":0.86}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
|
||||
|
||||
config = AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
skip_softmax=SkipSoftmaxSpec(
|
||||
target_sparsity=0.65,
|
||||
disabled_until_timestep=0.86,
|
||||
),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
Start with `target_sparsity=0.65` and `disabled_until_timestep=0.86`. Increase
|
||||
sparsity only after comparing quality against dense output at the same seed.
|
||||
The benefit grows with sequence length because only attention is accelerated.
|
||||
|
||||
## SAGE quantization
|
||||
|
||||
TRTLLM SAGE quantizes Q/K per block and V per channel so both attention
|
||||
matrix multiplications use low precision.
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
| --- | --- | --- |
|
||||
| `dtype_qk` | `int8`, `fp8_e4m3` | Q/K quantization dtype; absent means dense |
|
||||
| `q_block_size` | `1`, `4`, `16` | Q scale block size; default `1` |
|
||||
| `k_block_size` | `1`, `4`, `16` | K scale block size; default `16` |
|
||||
|
||||
V is always quantized per channel to FP8 E4M3 and K-smoothing is internal.
|
||||
Only block sizes with compiled kernels are accepted.
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
|
||||
"quant":{"dtype_qk":"fp8_e4m3","q_block_size":1,"k_block_size":16}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
|
||||
|
||||
config = AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
quant=AttnQuantSpec(
|
||||
dtype_qk="fp8_e4m3",
|
||||
q_block_size=1,
|
||||
k_block_size=16,
|
||||
),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
This path requires FlashInfer 0.6.16rc1 or newer. FP8 Q/K kernels are
|
||||
available on sm_100 and sm_103; INT8 Q/K kernels are compiled for sm_100 only.
|
||||
The shared `AttnQuantSpec` is backend typed: values intended for
|
||||
`FLASHINFER_ATTN` are rejected here and vice versa.
|
||||
@@ -1,389 +1,75 @@
|
||||
# CPU Offloading for Diffusion Models
|
||||
|
||||
## Overview
|
||||
vLLM-Omni provides three CPU-offload strategies for diffusion models. Choose
|
||||
the coarsest strategy that meets your memory target: finer-grained streaming
|
||||
saves more device memory but adds more transfer and setup overhead.
|
||||
|
||||
vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models:
|
||||
For the shared factory, discovery, and lifecycle contract, see the
|
||||
[CPU Offloading design](../../design/feature/offloader/README.md).
|
||||
|
||||
1. **Model-level (Sequential) Offloading**: Mutual exclusion between DiT model and encoder - only one is on GPU at a time.
|
||||
2. **Layerwise (Blockwise) Offloading**: Keeps only one transformer block on GPU at a time with compute-memory overlap.
|
||||
## Choose a strategy
|
||||
|
||||
Both strategies use pinned memory for faster CPU-GPU transfers. The strategies are **mutually exclusive** for now - if both are enabled, layerwise takes priority.
|
||||
| Strategy | Device residency | Parallel scope | Guide |
|
||||
| --- | --- | --- | --- |
|
||||
| Model-level (sequential) | One pipeline component group at a time | Single device | [Model-Level Offloading](offloader/cpu_offload.md) |
|
||||
| Layerwise (blockwise) | One transformer block, with next-block prefetch | Single device | [Layerwise Offloading](offloader/layerwise_offload.md) |
|
||||
| Distributed layerwise | Fixed two-block device buffer; optional host-weight sharding and AllGather | Multiple GPU/NPU ranks | [Distributed Layerwise Offloading](offloader/distributed_layerwise_offload.md) |
|
||||
|
||||
All strategies use pinned host memory for faster transfers where applicable.
|
||||
Configuration priority is:
|
||||
|
||||
1. Distributed layerwise offloading.
|
||||
2. Layerwise offloading.
|
||||
3. Model-level offloading.
|
||||
|
||||
Treat the flags as mutually exclusive. If more than one is enabled, the
|
||||
higher-priority strategy is selected.
|
||||
|
||||
## Quick selection
|
||||
|
||||
- Use [model-level offloading](offloader/cpu_offload.md) when swapping whole
|
||||
encoders and DiTs is enough to fit and phase-boundary transfers are
|
||||
acceptable.
|
||||
- Use [layerwise offloading](offloader/layerwise_offload.md) for compute-heavy video
|
||||
DiTs where block transfers can overlap computation.
|
||||
- Use [distributed layerwise offloading](offloader/distributed_layerwise_offload.md)
|
||||
when a multi-rank deployment also needs bounded device residency and,
|
||||
optionally, sharded host weights.
|
||||
|
||||
## Supported models
|
||||
|
||||
| Architecture | Example models | DiT class | Model-level | Layerwise | Distributed layerwise | Layerwise block attributes |
|
||||
| --- | --- | --- | :---: | :---: | :---: | --- |
|
||||
| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | yes | yes | — | `transformer_blocks`, `single_transformer_blocks` |
|
||||
| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | — | yes | — | `transformer_blocks`, `single_transformer_blocks` |
|
||||
| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | — | yes | — | `layers` |
|
||||
| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | — | yes | — | `transformer` |
|
||||
| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | yes | yes | — | `transformer_blocks` |
|
||||
| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | yes | yes | — | `down_blocks`, `up_blocks` |
|
||||
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | — | yes | — | `transformer_blocks` |
|
||||
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
|
||||
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
|
||||
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` | yes | yes | — | `layers` |
|
||||
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | — | yes | — | `layers`, customized modules |
|
||||
| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | yes | yes | yes | `layers`, `gen_layers` |
|
||||
|
||||
Model-level support requires discoverable DiT and encoder components.
|
||||
Layerwise support requires transformer block topology. Distributed support
|
||||
reuses that topology but still requires validation for the model, checkpoint,
|
||||
and parallel configuration.
|
||||
|
||||
## Compatibility anchors
|
||||
|
||||
These headings preserve links to sections that moved into dedicated guides.
|
||||
|
||||
## Model-level (Sequential) Offloading
|
||||
|
||||
### How It Works
|
||||
|
||||
Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks:
|
||||
|
||||
- **When encoders run**: DiT transformer is offloaded to CPU
|
||||
- **When DiT runs**: Encoders are offloaded to CPU, if more than one dit models, only one loaded on GPU, others get offloaded to CPU.
|
||||
- **VAE**: Stays resident on GPU
|
||||
|
||||
Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed.
|
||||
|
||||
### Usage
|
||||
|
||||
**Python API:**
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True)
|
||||
```
|
||||
|
||||
**CLI:**
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-cpu-offload
|
||||
```
|
||||
|
||||
### To Support a Model
|
||||
|
||||
Implement the `SupportsComponentDiscovery` protocol to declare which
|
||||
submodules serve as pipeline components (used by offloading, HSDP
|
||||
sharding, and other framework features):
|
||||
|
||||
```python
|
||||
from typing import ClassVar
|
||||
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
|
||||
|
||||
class MyPipeline(nn.Module, SupportsComponentDiscovery):
|
||||
_dit_modules: ClassVar[list[str]] = ["transformer"]
|
||||
_encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
|
||||
_vae_modules: ClassVar[list[str]] = ["vae"]
|
||||
_resident_modules: ClassVar[list[str]] = [] # optional
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.transformer = ... # DiT — stays on GPU during denoising
|
||||
self.text_encoder = ... # Encoder — offloaded to CPU during denoising
|
||||
self.vision_model = ... # Encoder — offloaded to CPU during denoising
|
||||
self.vae = ... # VAE — always on GPU
|
||||
```
|
||||
|
||||
- `_dit_modules`: attribute names of denoising submodules (kept on GPU
|
||||
during the diffusion loop).
|
||||
- `_encoder_modules`: attribute names of encoder/vision submodules
|
||||
(offloaded to CPU during the diffusion loop).
|
||||
- `_vae_modules`: attribute names of VAE(s) (always kept on GPU, not
|
||||
part of the mutual exclusion hooks).
|
||||
- `_resident_modules`: attribute names of small submodules that must
|
||||
stay on GPU during layerwise offloading (e.g. embedders, connectors).
|
||||
Optional — defaults to `[]`.
|
||||
|
||||
All attribute names support dotted paths for nested submodules
|
||||
(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).
|
||||
|
||||
Both DiT and encoder lists are needed because the offload hooks use
|
||||
mutual exclusion: when one group runs, the other moves to CPU.
|
||||
|
||||
### Limitations
|
||||
- Cold start latency increases
|
||||
- Adds overhead from CPU-GPU transfers between encoder and denoising phases
|
||||
- Support single GPU only for now
|
||||
|
||||
|
||||
### Component offloading for split models (e.g. Cosmos3)
|
||||
|
||||
Some models split their transformer into mutually-exclusive *components* that run
|
||||
in different phases of a single forward pass rather than as separate pipeline
|
||||
components -- e.g. Cosmos3's understanding (reasoner) component runs once per
|
||||
generation while the generation (generator) component runs every denoising step.
|
||||
Such models have no separate text encoder to swap against, so the transformer
|
||||
owns a small model-local offload path and wraps each phase with
|
||||
`with self._offload_context(name):`
|
||||
|
||||
```python
|
||||
class Cosmos3VFMTransformer(nn.Module):
|
||||
def forward(self, ...):
|
||||
with self._offload_context("reasoner"):
|
||||
... # understanding pass, runs once
|
||||
with self._offload_context("generator"):
|
||||
... # denoising pass, runs every step
|
||||
```
|
||||
|
||||
Model-level offloading then keeps exactly one component GPU-resident at a time
|
||||
(the other on CPU), reusing the same `SequentialOffloadHook` `.to()` movers. The
|
||||
pipeline opts in by exposing `enable_omni_model_cpu_offload` (which drives the
|
||||
transformer's `enable_model_cpu_offload` and pins the VAE). Layerwise offloading
|
||||
works for these models too -- each component declares its own block container via
|
||||
`_layerwise_offload_blocks_attrs`.
|
||||
|
||||
See [Model-Level Offloading](offloader/cpu_offload.md).
|
||||
|
||||
## Layerwise (Blockwise) Offloading
|
||||
|
||||
### How It Works
|
||||
|
||||
Layerwise offloading keeps only one transformer block on GPU at a time.
|
||||
|
||||
As each block completes, the next block is prefetched to GPU while the current block is freed. The pre and forward hooks utilized by layerwise offloading apply a separate CUDA stream (`copy_stream`) to overlap weight transfer with computation, and retain flattened tensors in pinned CPU memory for block parameters re-materialization. Encoders, VAE, and non-block DiT modules (embeddings, norms) always stay on GPU.
|
||||
|
||||
**Execution Flow:**
|
||||
|
||||
| Block | Pre-forward Hook | Forward | Post-forward Hook |
|
||||
|-------|------------------|---------|-------------------|
|
||||
| block-0 | Prefetch block-1 (async) | Compute block-0 | Free block-0 |
|
||||
| block-1 | Prefetch block-2 (async) | Compute block-1 | Free block-1 |
|
||||
| ... | ... | ... | ... |
|
||||
| block-(n-1) | **Prefetch block-0** (async) | Compute block-(n-1) | Free block-(n-1) |
|
||||
|
||||
Each transformer block has a `LayerwiseOffloadHook` that prefetches the next block before forward and frees the current block after forward.
|
||||
|
||||
Layerwise offloading is primarily recommended for large **video generation models** where the compute cost per block is high enough to effectively overlap with memory prefetch operations. For example, Wan2.2 T2V and I2V pipelines.
|
||||
|
||||
### Usage
|
||||
|
||||
**Python API:**
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
# Text-to-video
|
||||
m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_layerwise_offload=True)
|
||||
|
||||
# Or image-to-video
|
||||
m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True)
|
||||
```
|
||||
|
||||
**CLI:**
|
||||
```bash
|
||||
# Text-to-video
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-layerwise-offload
|
||||
|
||||
# Or image-to-video
|
||||
vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --enable-layerwise-offload
|
||||
```
|
||||
|
||||
### To Support a Model
|
||||
|
||||
Models must define the blocks attribute name for layerwise offloading:
|
||||
|
||||
```python
|
||||
class WanTransformer3DModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = ["blocks"] # Attribute names containing transformer blocks
|
||||
|
||||
def __init__(self):
|
||||
self.blocks = nn.ModuleList([...]) # Transformer blocks
|
||||
```
|
||||
|
||||
For models with multiple block types:
|
||||
|
||||
```python
|
||||
class Flux2Transformer2DModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
|
||||
```
|
||||
|
||||
### Limitations
|
||||
- Cold start latency increases because offloaded components must be moved to CPU
|
||||
during setup; layerwise offload may add extra weight consolidation and pinning
|
||||
work.
|
||||
- Performance depends on compute cost and H2D bandwidth as well
|
||||
- Support single GPU only for now
|
||||
|
||||
See [Layerwise Offloading](offloader/layerwise_offload.md).
|
||||
|
||||
## Distributed Layerwise Offloading
|
||||
|
||||
### How It Works
|
||||
|
||||
Distributed layerwise offloading extends single-GPU layerwise offloading to
|
||||
multi-device deployments. In the default mode, each rank stores only
|
||||
**1/dp_size** of the model weights in host memory; full layer weights are
|
||||
reconstructed at runtime via **AllGather** on a dedicated communication stream,
|
||||
overlapped with computation via a fixed double-buffer scheme. The
|
||||
`--dlo-no-use-allgather` mode instead streams standard-loader rank-local
|
||||
weights independently and does not shard weights across ranks.
|
||||
|
||||
For the implementation design and compatibility matrix, see
|
||||
[Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md).
|
||||
|
||||
**Key features:**
|
||||
- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights,
|
||||
reconstructed per-layer via `all_gather_into_tensor`
|
||||
- **Fixed double-buffer**: exactly 2 transformer blocks on each device at any
|
||||
time, regardless of model size
|
||||
- **DP multi-concurrency**: N concurrent requests processed in parallel
|
||||
(AllGather only gathers weight shards, which are request-independent)
|
||||
- **mmap weight loading**: weights loaded as mmap views pointing to shared OS
|
||||
page cache, eliminating O(dp_size × model_size) RSS during model creation
|
||||
- **Platform-agnostic**: works on NVIDIA GPU (CUDA/NCCL) and Ascend NPU
|
||||
(CANN/HCCL) via vLLM-Omni's platform abstraction
|
||||
|
||||
**Execution Flow (per device):**
|
||||
|
||||
```
|
||||
Compute Stream: [Layer N] [Layer N+1] [Layer N+2]
|
||||
H2D Stream: [H2D Shard N+1] [H2D Shard N+2]
|
||||
AllGather: [AG N+1] [AG N+2]
|
||||
Slot Usage: Slot0: Layer N Slot1: Layer N+1
|
||||
```
|
||||
|
||||
### Usage
|
||||
|
||||
**CLI:**
|
||||
```bash
|
||||
# 4× GPU/NPU with AllGather (recommended)
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4
|
||||
|
||||
# Without DLO AllGather (each rank streams standard-loader rank-local weights)
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4 \
|
||||
--dlo-no-use-allgather
|
||||
|
||||
# With SP instead of DP (long sequences)
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--usp 4
|
||||
```
|
||||
|
||||
**Python API:**
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
m = Omni(
|
||||
model="/path/to/model",
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True, # default
|
||||
)
|
||||
```
|
||||
|
||||
### CLI Flags
|
||||
|
||||
| Flag | Description | Default |
|
||||
|------|-------------|---------|
|
||||
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
|
||||
| `--data-parallel-size N` | Number of DP ranks; DLO weight-sharding group when AllGather is enabled | `1` |
|
||||
| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` |
|
||||
| `--dlo-no-use-allgather` | Use standard-loader rank-local weights; no DLO sharding or AllGather | `false` |
|
||||
|
||||
### How model weights are loaded (mmap path)
|
||||
|
||||
When DLO + AllGather is active, the offloader:
|
||||
|
||||
1. **Saves non-persistent buffers** (e.g. RoPE `inv_freq`, timestep `freqs`)
|
||||
from the normally-created transformer
|
||||
2. **Converts to meta device** via `to_empty(device="meta")`, releasing random
|
||||
initialization weights
|
||||
3. **Loads checkpoint weights as mmap views** via `safe_open().get_tensor()`,
|
||||
which return views into the OS page cache (shared across all ranks, 0 RSS)
|
||||
4. **Calls `post_load_weights()`** to apply model-specific dtype conversions
|
||||
(e.g. Cosmos3's `time_embedder` → FP32)
|
||||
5. **Restores non-persistent buffers** from saved copies
|
||||
|
||||
This approach requires **zero model-specific code changes** — no pipeline or
|
||||
transformer modifications are needed.
|
||||
|
||||
### OffloadPlan (declarative topology metadata)
|
||||
|
||||
Models can optionally declare an `OffloadPlan` class variable to provide
|
||||
topology metadata (block attribute names, submodules to offload) without
|
||||
any offload-specific logic:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.offloader import OffloadPlan
|
||||
|
||||
class MyPipeline(nn.Module):
|
||||
_dit_modules = ["transformer"]
|
||||
_offload_plan = OffloadPlan(
|
||||
block_attrs={"transformer": ("blocks",)},
|
||||
offload_submodules={"context_encoder": "layers"},
|
||||
)
|
||||
```
|
||||
|
||||
When not declared, the offloader falls back to `_layerwise_offload_blocks_attrs`
|
||||
and heuristic attribute search. This is backward-compatible — existing models
|
||||
work without any changes.
|
||||
|
||||
### DP Multi-concurrency
|
||||
|
||||
When `--data-parallel-size > 1` and AllGather is enabled, the scheduler batches
|
||||
up to `dp_size` requests per denoise step. Each DP rank processes a different
|
||||
request while AllGather synchronizes only weight shards (request-independent).
|
||||
|
||||
**Requirements for concurrent requests:**
|
||||
- `num_inference_steps` must be specified explicitly (None is not allowed)
|
||||
- All concurrent requests must have the same `num_inference_steps` value
|
||||
- These constraints exist because AllGather is a collective that requires
|
||||
every rank to participate at each denoise step
|
||||
|
||||
### Limitations
|
||||
- Online quantization (FP8) is incompatible with mmap loading — falls back to
|
||||
regular loading only when the model is outside the DLO+AllGather mmap path;
|
||||
online quantization with DLO+AllGather is rejected. Use
|
||||
`--dlo-no-use-allgather` or disable online quantization.
|
||||
- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware
|
||||
weight-loader callbacks are bypassed. The no-AllGather path retains the
|
||||
standard TP loader and is experimental.
|
||||
- HSDP + AllGather is rejected because it would double-shard parameters. HSDP
|
||||
with no-AllGather is accepted at configuration level, but has limited
|
||||
end-to-end validation.
|
||||
- `num_inference_steps=None` is not allowed in DP multi-concurrency mode
|
||||
|
||||
**Module Discovery**
|
||||
|
||||
The offloader discovers pipeline components in two ways:
|
||||
|
||||
1. **Protocol-based** (preferred): If the pipeline implements
|
||||
`SupportsComponentDiscovery`, its `_dit_modules`, `_encoder_modules`,
|
||||
`_vae_modules`, and `_resident_modules` class variables are used
|
||||
directly. All attribute names support dotted paths (e.g.
|
||||
`"pipe.transformer"`, `"bagel.time_embedder"`) for nested submodules.
|
||||
|
||||
2. **Fallback attribute scan**: Otherwise, the offloader scans for
|
||||
well-known attribute names:
|
||||
- **DiT modules**: `transformer`, `transformer_2`, `dit`, `sr_dit`, `language_model`, `transformer_blocks`, `model`
|
||||
- **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder`
|
||||
- **VAE**: `vae`, `audio_vae`
|
||||
|
||||
**Hook System**
|
||||
|
||||
Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelHook`) to register pre/post forward callbacks on modules, enabling automatic swapping without modifying model code.
|
||||
|
||||
**Backend Architecture**
|
||||
|
||||
```
|
||||
OffloadBackend (base class)
|
||||
├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap)
|
||||
│ (delegates to a pipeline's enable_omni_model_cpu_offload
|
||||
│ for split models like Cosmos3)
|
||||
├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
|
||||
│ (single-GPU, full weights on host)
|
||||
└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
|
||||
(multi-GPU, sharded AllGather or rank-local mode)
|
||||
```
|
||||
|
||||
Factory function `get_offload_backend()` selects the appropriate backend based on
|
||||
configuration.
|
||||
|
||||
For split models, `ModelLevelOffloadBackend.enable()` detects a pipeline's
|
||||
`enable_omni_model_cpu_offload` hook and delegates to it; Cosmos3 then swaps its
|
||||
reasoner/generator components inside the model forward pass.
|
||||
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Distributed Layerwise Offload | Blocks Attrs (Layerwise specific) |
|
||||
|--------------|----------------|-----------|---------------------|-------------------|-------------------------------|-----------------------------------|
|
||||
| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
|
||||
| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
|
||||
| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | - | `"layers"` |
|
||||
| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | - | `"transformer"` |
|
||||
| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"` |
|
||||
| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | - | `"down_blocks"`, `"up_blocks"` |
|
||||
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | - | `"transformer_blocks"` |
|
||||
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
|
||||
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
|
||||
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | - | `"layers"` |
|
||||
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | - | `"layers"`, `"customized modules"` |
|
||||
| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | ✓ | ✓ | ✓ | `"layers"`, `"gen_layers"` |
|
||||
|
||||
**Notes:**
|
||||
- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
|
||||
- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
|
||||
- Distributed Layerwise Offloading reuses the layerwise block topology, but the
|
||||
model and parallelism combination still needs validation. AllGather mode
|
||||
additionally requires a compatible checkpoint/mmap layout; no-AllGather mode
|
||||
retains the standard loader's rank-local tensors. See the
|
||||
[DLO design](../../design/feature/distributed_layerwise_offload.md) and
|
||||
[Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
|
||||
for usage examples.
|
||||
See [Distributed Layerwise Offloading](offloader/distributed_layerwise_offload.md).
|
||||
@@ -0,0 +1,81 @@
|
||||
# Model-Level Offloading
|
||||
|
||||
Model-level, or sequential, offloading keeps only the pipeline component group
|
||||
currently executing on the accelerator. It is the simplest offload strategy
|
||||
and is selected with `--enable-cpu-offload`.
|
||||
|
||||
## How it works
|
||||
|
||||
Pre-forward hooks enforce mutual exclusion between DiT and encoder modules:
|
||||
|
||||
- before an encoder runs, the DiT moves to CPU;
|
||||
- before a DiT runs, encoders and other DiTs move to CPU; and
|
||||
- VAE modules remain on the accelerator.
|
||||
|
||||
Pinned host memory reduces transfer overhead. Transfers occur at phase
|
||||
boundaries, so cold-start and encoder-to-denoiser transitions become slower.
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
enable_cpu_offload=True,
|
||||
)
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--omni --enable-cpu-offload
|
||||
```
|
||||
|
||||
## Model integration
|
||||
|
||||
Pipelines should implement `SupportsComponentDiscovery`:
|
||||
|
||||
```python
|
||||
from typing import ClassVar
|
||||
|
||||
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
|
||||
|
||||
|
||||
class MyPipeline(nn.Module, SupportsComponentDiscovery):
|
||||
_dit_modules: ClassVar[list[str]] = ["transformer"]
|
||||
_encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
|
||||
_vae_modules: ClassVar[list[str]] = ["vae"]
|
||||
_resident_modules: ClassVar[list[str]] = []
|
||||
```
|
||||
|
||||
All entries may be dotted paths. DiT and encoder lists are both required for
|
||||
mutual exclusion. VAE modules are pinned but not swapped; resident modules are
|
||||
small modules that must stay on the accelerator for layerwise paths.
|
||||
|
||||
## Split-model components
|
||||
|
||||
Some models, such as Cosmos3, split one transformer into mutually exclusive
|
||||
components that run in different phases. The pipeline exposes
|
||||
`enable_omni_model_cpu_offload`, and the backend delegates to the model-local
|
||||
contexts:
|
||||
|
||||
```python
|
||||
class Cosmos3VFMTransformer(nn.Module):
|
||||
def forward(self, ...):
|
||||
with self._offload_context("reasoner"):
|
||||
...
|
||||
with self._offload_context("generator"):
|
||||
...
|
||||
```
|
||||
|
||||
This preserves the same invariant—exactly one component is device resident—
|
||||
while reusing sequential `.to()` movers.
|
||||
|
||||
## Limitations
|
||||
|
||||
- Single device only.
|
||||
- Higher cold-start latency.
|
||||
- Transfers between encoder and denoising phases add latency.
|
||||
|
||||
See the [model-level design](../../../design/feature/offloader/cpu_offload.md)
|
||||
for lifecycle and extension invariants.
|
||||
@@ -0,0 +1,180 @@
|
||||
# Distributed Layerwise Offloading
|
||||
|
||||
Distributed layerwise offloading (DLO) extends block streaming to multi-device
|
||||
deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
|
||||
the host weights and reconstructs each layer at runtime. Without AllGather,
|
||||
each rank streams a complete block independently. Compatible TP1 deployments
|
||||
can share checkpoint-backed host pages among processes on the same node;
|
||||
otherwise DLO streams the ordinary loader's rank-local tensors.
|
||||
|
||||
See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
|
||||
for the implementation contract and compatibility matrix.
|
||||
|
||||
## Execution model
|
||||
|
||||
DLO overlaps three operations with a fixed two-block device buffer:
|
||||
|
||||
```text
|
||||
Compute stream: [Layer N] [Layer N+1] [Layer N+2]
|
||||
H2D stream: [H2D shard N+1] [H2D shard N+2]
|
||||
AllGather: [AG N+1] [AG N+2]
|
||||
Slots: slot 0: Layer N slot 1: Layer N+1
|
||||
```
|
||||
|
||||
AllGather communicates only request-independent weight shards, so data-
|
||||
parallel ranks may process different requests concurrently.
|
||||
|
||||
## Usage
|
||||
|
||||
```bash
|
||||
# Four ranks with sharded host weights and AllGather
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4
|
||||
|
||||
# Standard-loader rank-local weights, without DLO AllGather
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4 \
|
||||
--dlo-no-use-allgather
|
||||
|
||||
# Sequence parallel deployment
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--usp 4
|
||||
|
||||
# Opt-in normalized online-FP8 cache (TP1 + DLO AllGather)
|
||||
vllm serve /path/to/model --omni \
|
||||
--quantization fp8 \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--dlo-fp8-cache-dir /path/to/node-local/fp8-cache
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="/path/to/model",
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
)
|
||||
```
|
||||
|
||||
## Flags
|
||||
|
||||
| Flag | Meaning | Default |
|
||||
| --- | --- | --- |
|
||||
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
|
||||
| `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
|
||||
| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
|
||||
| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
|
||||
| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
|
||||
| `--dlo-fp8-cache-dir PATH` | Opt into the Phase-I normalized online-FP8 cache for the transformer | unset |
|
||||
|
||||
## Host-weight loading
|
||||
|
||||
The diffusion loader chooses host storage before DLO is enabled. It first
|
||||
attempts to build a complete, validated direct-checkpoint mmap plan. If names,
|
||||
coverage, shape, dtype, topology, or loader-callback compatibility cannot be
|
||||
proven, it runs the ordinary model loader instead. DLO consumes that result and
|
||||
does not make a second checkpoint-compatibility decision.
|
||||
|
||||
The shared-mmap optimization in this phase is supported only with TP1. TP
|
||||
greater than one falls back before model mutation to ordinary TP-aware loading.
|
||||
DLO may still consume those TP-local tensors, but this is a compatibility path:
|
||||
it does not share checkpoint-backed runtime weights across DP replicas and
|
||||
provides no shared-mmap host-memory guarantee.
|
||||
|
||||
The mmap plan skips only dedicated DiT weight sources. Other component sources,
|
||||
such as a text encoder loaded through the shared diffusion loader, continue to
|
||||
use their ordinary component loader. A checkpoint source that mixes DiT and
|
||||
non-DiT weights falls back completely rather than leaving an unplanned
|
||||
component uninitialized.
|
||||
|
||||
With direct checkpoint mmap, the loader:
|
||||
|
||||
1. saves non-persistent buffers such as RoPE frequencies;
|
||||
2. moves the normally created transformer to the meta device;
|
||||
3. loads checkpoint tensors as mmap views backed by the shared OS page cache;
|
||||
4. applies any loader-owned bounded layout adapters while packing blocks;
|
||||
5. restores saved non-persistent buffers; and
|
||||
6. preserves `post_load_weights()` and `validate_loaded_weights()` lifecycle
|
||||
hooks.
|
||||
|
||||
For AllGather with a group larger than one, each process copies only its
|
||||
persistent shard and then releases the source mapping. For no-AllGather, each
|
||||
process keeps the mapping open and packs complete blocks through two bounded
|
||||
pinned staging slots. Processes mapping the same files on one node share the
|
||||
immutable pages; no-AllGather still performs a complete-block H2D copy in each
|
||||
process.
|
||||
|
||||
When the effective DLO group size is one, `dlo_use_allgather=True` does not
|
||||
perform a collective and uses the same rank-local transfer behavior.
|
||||
|
||||
## Normalized online-FP8 cache
|
||||
|
||||
`--dlo-fp8-cache-dir` is an opt-in Phase-I cache for per-tensor online FP8 with
|
||||
TP1 and DLO AllGather. On a cache miss, the ordinary loader quantizes the
|
||||
transformer once and publishes finalized FP8 weights and scales as an atomic,
|
||||
serialized safetensors artifact. On a valid cache hit, the loader selects that
|
||||
serialized source before DLO planning, so the transformer can be loaded through
|
||||
the direct checkpoint-mmap path without re-materializing the full BF16 source.
|
||||
|
||||
The cache is keyed to the source files through the manifest's schema, component,
|
||||
quantization contract, and source fingerprint. Incomplete, incompatible, or
|
||||
stale artifacts are treated as cache misses. Cache creation is best-effort: a
|
||||
publication failure logs a warning and keeps the ordinary online-FP8 path.
|
||||
|
||||
Phase I normalizes the transformer only. Other pipeline components keep their
|
||||
ordinary loading and quantization behavior. The cache is not selected when DLO
|
||||
is disabled, AllGather is disabled, TP is greater than one, or HSDP is enabled.
|
||||
|
||||
## Declarative topology
|
||||
|
||||
Models may declare an `OffloadPlan` instead of embedding offload logic:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.offloader import OffloadPlan
|
||||
|
||||
|
||||
class MyPipeline(nn.Module):
|
||||
_dit_modules = ["transformer"]
|
||||
_offload_plan = OffloadPlan(
|
||||
block_attrs={"transformer": ("blocks",)},
|
||||
offload_submodules={"context_encoder": "layers"},
|
||||
)
|
||||
```
|
||||
|
||||
When no plan exists, discovery falls back to
|
||||
`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup.
|
||||
|
||||
## Data-parallel concurrency
|
||||
|
||||
With `data_parallel_size > 1` and AllGather enabled, the scheduler can process
|
||||
up to `dp_size` requests per denoising step. Every concurrent request must set
|
||||
the same explicit `num_inference_steps`; `None` is rejected because every rank
|
||||
must enter each collective.
|
||||
|
||||
## Limitations
|
||||
|
||||
- Direct checkpoint mmap currently requires TP1. TP greater than one is
|
||||
outside the Phase A shared-mmap support scope and falls back before model
|
||||
mutation to the ordinary TP-aware loader. DLO can stream that runtime layout,
|
||||
but it provides no shared-mmap host-memory benefit or guarantee.
|
||||
- HSDP plus AllGather is rejected to avoid double sharding. HSDP without
|
||||
AllGather has limited end-to-end validation.
|
||||
- Per-tensor online FP8 linears use the ordinary loader and can run with either
|
||||
DLO transfer path. With AllGather, every rank temporarily materializes the
|
||||
complete FP8 model in host memory before DLO retains only its shard. Other
|
||||
online quantization methods require no-AllGather until their runtime layouts
|
||||
are validated.
|
||||
- Resident leading layers require `--dlo-no-use-allgather` and a model
|
||||
`OffloadPlan` that declares eligible `resident_dit_paths`.
|
||||
- DP concurrency requires an explicit, identical inference-step count.
|
||||
|
||||
Sharing transformed TP or quantized runtime layouts beyond this TP1 transformer
|
||||
cache remains a follow-up design in
|
||||
[RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
|
||||
|
||||
See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
|
||||
for an end-to-end example.
|
||||
@@ -0,0 +1,65 @@
|
||||
# Layerwise Offloading
|
||||
|
||||
Layerwise, or blockwise, offloading keeps one transformer block on the
|
||||
accelerator and prefetches the next block while the current block computes.
|
||||
It is best suited to compute-heavy video DiTs whose block execution can hide
|
||||
host-to-device transfers.
|
||||
|
||||
## Execution flow
|
||||
|
||||
Each block has a pre-forward and post-forward hook. Parameters are consolidated
|
||||
in pinned host tensors and rematerialized for execution on a dedicated copy
|
||||
stream.
|
||||
|
||||
| Block | Pre-forward hook | Forward | Post-forward hook |
|
||||
| --- | --- | --- | --- |
|
||||
| block 0 | Prefetch block 1 | Compute block 0 | Free block 0 |
|
||||
| block 1 | Prefetch block 2 | Compute block 1 | Free block 1 |
|
||||
| ... | ... | ... | ... |
|
||||
| last block | Prefetch block 0 | Compute last block | Free last block |
|
||||
|
||||
Encoders, VAE modules, and non-block DiT modules such as embeddings and norms
|
||||
remain device resident.
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
enable_layerwise_offload=True,
|
||||
)
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--omni --enable-layerwise-offload
|
||||
```
|
||||
|
||||
## Model integration
|
||||
|
||||
Transformer classes declare containers of executable blocks:
|
||||
|
||||
```python
|
||||
class WanTransformer3DModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = ["blocks"]
|
||||
|
||||
|
||||
class Flux2Transformer2DModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = [
|
||||
"transformer_blocks",
|
||||
"single_transformer_blocks",
|
||||
]
|
||||
```
|
||||
|
||||
See the [layerwise design](../../../design/feature/offloader/layerwise_offload.md) for
|
||||
the discovery and hook invariants. `OffloadPlan` is a separate declarative
|
||||
topology path for distributed layerwise offload.
|
||||
|
||||
## Limitations
|
||||
|
||||
- Single device only.
|
||||
- Setup consolidates and pins block parameters, increasing cold-start time.
|
||||
- Performance depends on block compute time and host-to-device bandwidth;
|
||||
lightweight blocks may not hide transfers.
|
||||
@@ -64,6 +64,41 @@ omni = Omni(
|
||||
)
|
||||
```
|
||||
|
||||
### Combined with Online FP8 Quantization
|
||||
|
||||
HSDP can be combined with online FP8 quantization for diffusion models that
|
||||
support both features. HSDP shards the quantized weights across GPUs to reduce
|
||||
the per-GPU weight memory footprint.
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
|
||||
omni = Omni(
|
||||
model="<your-model>",
|
||||
quantization="fp8",
|
||||
parallel_config=DiffusionParallelConfig(
|
||||
use_hsdp=True,
|
||||
hsdp_shard_size=2,
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
For online serving:
|
||||
|
||||
```bash
|
||||
vllm serve <your-model> \
|
||||
--omni \
|
||||
--quantization fp8 \
|
||||
--use-hsdp \
|
||||
--hsdp-shard-size 2
|
||||
```
|
||||
|
||||
The model must support both HSDP and online FP8 quantization. HSDP remains
|
||||
incompatible with Tensor Parallelism. See the
|
||||
[FP8 Quantization Guide](../../quantization/fp8.md) for supported models and
|
||||
hardware.
|
||||
|
||||
---
|
||||
|
||||
## Example Script
|
||||
@@ -117,7 +152,8 @@ In `DiffusionParallelConfig`:
|
||||
**Constraints:**
|
||||
|
||||
- `hsdp_replicate_size × hsdp_shard_size == world_size`
|
||||
- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
|
||||
- HSDP cannot be used with Tensor Parallelism, ordinary Data Parallelism,
|
||||
Pipeline Parallelism, or Expert Parallelism.
|
||||
|
||||
---
|
||||
|
||||
@@ -147,3 +183,4 @@ For detailed instructions on adding HSDP support to new models, see the [HSDP Co
|
||||
1. ✅ **Enable HSDP** - Set `use_hsdp=True` and `hsdp_shard_size` to reduce per-GPU memory for large models
|
||||
2. ✅ **Combine with SP** - Use together with `ulysses_degree` for video models requiring both memory reduction and sequence parallelism
|
||||
3. ⚠️ **Incompatible with TP** - `tensor_parallel_size` must be 1 when HSDP is enabled
|
||||
4. ✅ **Combine with online FP8** - Use `quantization="fp8"` or `--quantization fp8` to reduce the sharded weight memory footprint
|
||||
@@ -106,7 +106,10 @@ The following tables show which models support each feature:
|
||||
|
||||
> Notes:
|
||||
|
||||
> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload.md).
|
||||
> 1. CPU Offload has three strategies: model-level, layerwise, and distributed
|
||||
> layerwise. The tables below show **layerwise support** only. Split models
|
||||
> like Cosmos3 swap their reasoner/generator components for model-level
|
||||
> offload; see the [CPU Offload Guide](diffusion/cpu_offload.md).
|
||||
> 2. The **💾Quantization** column is collapsed for readability. See [Quantization](quantization/overview.md) for per-method and per-model support details.
|
||||
|
||||
### ImageGen
|
||||
@@ -158,6 +161,7 @@ The following tables show which models support each feature:
|
||||
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2.5** | ❌ | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ |
|
||||
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
@@ -171,7 +175,7 @@ The following tables show which models support each feature:
|
||||
**Frame Interpolation Support**
|
||||
|
||||
- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
|
||||
- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, Helios, HunyuanVideo-1.5, DreamID-Omni
|
||||
- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni
|
||||
|
||||
### AudioGen
|
||||
|
||||
@@ -214,7 +218,7 @@ The following tables show which models support each feature:
|
||||
3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
|
||||
4. The CPU Offloading (Layerwise) row describes local layerwise offload.
|
||||
Multi-device Distributed Layerwise Offload has a separate topology and
|
||||
compatibility matrix in the [CPU Offloading guide](diffusion/cpu_offload.md#distributed-layerwise-offloading).
|
||||
compatibility matrix in the [Distributed Layerwise Offloading guide](diffusion/offloader/distributed_layerwise_offload.md).
|
||||
5. The compatibility matrix uses FP8 as the representative quantization method.
|
||||
6. Step Execution is not compatible with any diffusion cache backend. LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).
|
||||
|
||||
|
||||
@@ -43,7 +43,7 @@ if __name__ == "__main__":
|
||||
),
|
||||
],
|
||||
)
|
||||
outputs[0].request_output.images[0].save("output.png")
|
||||
outputs[0].images[0].save("output.png")
|
||||
```
|
||||
|
||||
## Image-to-Image (Image Editing)
|
||||
@@ -72,7 +72,7 @@ if __name__ == "__main__":
|
||||
),
|
||||
],
|
||||
)
|
||||
outputs[0].request_output.images[0].save("output.png")
|
||||
outputs[0].images[0].save("output.png")
|
||||
```
|
||||
|
||||
## Generation Parameters
|
||||
|
||||
@@ -55,7 +55,7 @@ if __name__ == "__main__":
|
||||
omni = Omni(model="Qwen/Qwen-Image")
|
||||
prompt = "a cup of coffee on the table"
|
||||
outputs = omni.generate(prompt)
|
||||
images = outputs[0].request_output.images
|
||||
images = outputs[0].images
|
||||
images[0].save("coffee.png")
|
||||
```
|
||||
|
||||
@@ -179,7 +179,7 @@ if __name__ == "__main__":
|
||||
]
|
||||
outputs = omni.generate(prompts)
|
||||
for i, output in enumerate(outputs):
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
!!! info
|
||||
@@ -207,7 +207,7 @@ if __name__ == "__main__":
|
||||
}
|
||||
])
|
||||
for i, output in enumerate(outputs):
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
You can also pass a negative prompt via the CLI argument `--negative-prompt`:
|
||||
|
||||
@@ -65,7 +65,7 @@ More failure modes (for example, network jitter and network interruption) will b
|
||||
|
||||
## Source of Conclusions
|
||||
|
||||
The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`), and **`VoxCPM2`** (`/v1/audio/speech`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`, `test_reliability_voxcpm2.py`) and runs weekly via `.buildkite/cuda/test-weekly.yml`.
|
||||
The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, and **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`) and runs weekly via `.buildkite/cuda/test-weekly.yml`.
|
||||
|
||||
### Example: `SIGTERM` fault injection log (Qwen3-Omni)
|
||||
|
||||
|
||||
@@ -165,7 +165,7 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \
|
||||
|
||||
- **Tensor Parallelism — Text Encoder Not Sharded**: TP currently only shards the DiT blocks. Each TP rank retains a **full copy of the text encoder weights**, leading to significant GPU memory overhead proportional to TP degree. Tracked in [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).
|
||||
|
||||
- **CPU Offloading — Two Modes Are Mutually Exclusive**: Model-level offload (`enable_cpu_offload`) and layerwise offload (`enable_layerwise_offload`) cannot be used simultaneously. If both are set, layerwise takes priority and model-level is silently ignored.
|
||||
- **CPU Offloading — Strategies Are Mutually Exclusive**: Distributed layerwise, layerwise, and model-level offload should not be enabled together. Selection priority is distributed layerwise, then layerwise, then model-level.
|
||||
|
||||
- **CPU Offloading — VAE stays on GPU**: Both offloading strategies keep the VAE on GPU at all times. For high-resolution generation, VAE decode can still cause OOM. Mitigate by combining with `vae_use_tiling=True` or VAE Patch Parallelism.
|
||||
|
||||
|
||||
@@ -80,12 +80,14 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
### Diffusion Models
|
||||
|
||||
| Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization |
|
||||
|-------|-----------|:-------:|:------:|----------------|------------------|------------------|
|
||||
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | |
|
||||
| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | |
|
||||
| LTX-2 | `Lightricks/LTX-2`, `rootonchair/LTX-2-19b-distilled` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
|
||||
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers`, `diffusers/LTX-2.3-Distilled-Diffusers` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
|
||||
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None | ✅︎ |
|
||||
| FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None | |
|
||||
| FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | Yes | Not validated | All layers | None | ✅︎ |
|
||||
@@ -93,7 +95,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
|
||||
| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
|
||||
| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |
|
||||
| Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | |
|
||||
| MiniMax H3 | `MiniMaxAI/MiniMax-H3` | Yes | Not validated | DiT linears except patch, timestep, and final projections; not compatible with layerwise offload | None | |
|
||||
| MiniMax-H3 | `MiniMaxAI/MiniMax-H3` (`FL2VA` / `Ref2VA`) | Yes | Not validated | `quantization="fp8"` quantizes eligible DiT and text-encoder linears; mixed-precision input/output heads stay FP32 | None | ✅︎ |
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
@@ -136,6 +138,39 @@ outputs = omni.generate(
|
||||
)
|
||||
```
|
||||
|
||||
### Global and per-component scope
|
||||
|
||||
A plain, global quantization configuration is passed unchanged to every
|
||||
quantization-aware component constructed by a pipeline. This includes an
|
||||
eligible encoder when that encoder is implemented with vLLM quantizable
|
||||
layers; it does not rewrite arbitrary `torch.nn` modules. A structured
|
||||
component map is the only way to narrow that scope. Pipeline integrations that
|
||||
do not yet expose an encoder through the quantization factory remain DiT-only.
|
||||
|
||||
For a pipeline that exposes both a transformer and a quantization-aware text
|
||||
encoder, the scope is:
|
||||
|
||||
| Configuration | Transformer | Text encoder | Components without supported quantizable layers |
|
||||
|---------------|-------------|--------------|-------------------------------------------------|
|
||||
| `quantization="fp8"` | FP8 | FP8 | checkpoint precision |
|
||||
| `{"transformer": {"method": "fp8"}}` | FP8 | checkpoint precision | checkpoint precision |
|
||||
| `{"text_encoder": {"method": "fp8"}}` | checkpoint precision | FP8 | checkpoint precision |
|
||||
|
||||
Use `quantization_config` for component-selective Python configuration. For
|
||||
example, quantize only the text decoder:
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="<your-model>",
|
||||
quantization_config={"text_encoder": {"method": "fp8"}},
|
||||
)
|
||||
```
|
||||
|
||||
Component keys are runtime prefixes exposed by the pipeline integration; common
|
||||
keys include `transformer` and `text_encoder`. Entries may be combined, and
|
||||
`ignored_layers` can keep named eligible layers in checkpoint precision. Check
|
||||
the model recipe for supported components and their runtime prefixes.
|
||||
|
||||
CLI:
|
||||
|
||||
```bash
|
||||
@@ -144,6 +179,24 @@ python text_to_image.py --model <your-model> --quantization fp8 --ignored-layers
|
||||
vllm serve <your-model> --omni --quantization fp8
|
||||
```
|
||||
|
||||
### Distributed layerwise offload
|
||||
|
||||
Per-tensor online FP8 linears can be combined with DLO's default AllGather path:
|
||||
|
||||
```bash
|
||||
vllm serve <your-model> \
|
||||
--omni \
|
||||
--quantization fp8 \
|
||||
--enable-distributed-layerwise-offload
|
||||
```
|
||||
|
||||
Direct checkpoint mmap is not used for this combination. Each rank runs the
|
||||
ordinary loader to create the final FP8 weights and scales before DLO shards
|
||||
them, so startup temporarily materializes the complete FP8 model in host memory
|
||||
on every rank. Runtime DLO host residency is sharded normally. Use
|
||||
`--dlo-no-use-allgather` when the synchronized AllGather request-wave contract
|
||||
does not fit the workload.
|
||||
|
||||
## Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|
||||
@@ -24,7 +24,6 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.model_executor.stage_input_processors.aura_omni import (
|
||||
DEFAULT_QWEN3_TTS_REF_TEXT,
|
||||
@@ -228,7 +227,6 @@ def parse_args():
|
||||
parser.add_argument("--init-timeout", type=int, default=2000, help="Overall initialization timeout.")
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=2000, help="Per-stage initialization timeout.")
|
||||
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.")
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -59,7 +59,7 @@ def main(args):
|
||||
|
||||
omni = Omni(
|
||||
model=args.model_name,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
deploy_config=args.deploy_config,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
batch_timeout=args.batch_timeout,
|
||||
@@ -167,10 +167,10 @@ def parse_args():
|
||||
help="Sampling rate for audio loading (default: 16000).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to stage configs YAML file.",
|
||||
help="Path to a deploy config YAML file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
|
||||
@@ -0,0 +1,288 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Generate a LingBot-World v2 video from an image and camera trajectory.
|
||||
|
||||
The official checkpoint is licensed separately under CC BY-NC-SA and is
|
||||
restricted to non-commercial use. This vLLM-Omni integration remains
|
||||
Apache-2.0 licensed.
|
||||
|
||||
Example:
|
||||
python examples/offline_inference/diffusion/lingbot_world_v2.py \
|
||||
--prompt "The camera moves slowly forward through the scene." \
|
||||
--image /path/to/first_frame.png \
|
||||
--action-dir /path/to/actions/forward \
|
||||
--output lingbot_world_v2.mp4
|
||||
|
||||
``--action-dir`` must contain ``poses.npy`` and ``intrinsics.npy``. The
|
||||
example configures its parent as the trusted action root and sends the
|
||||
contained directory through ``extra_args["action_path"]``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import math
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
|
||||
_NUM_INFERENCE_STEPS = 4
|
||||
_MAX_SEQUENCE_LENGTH = 512
|
||||
_MAX_PIXEL_AREA = 480 * 832
|
||||
_MAX_RAW_FRAMES = 117
|
||||
_MAX_SOURCE_CAMERA_FRAMES = 4096
|
||||
_TEMPORAL_COMPRESSION = 4
|
||||
_LATENT_FRAMES_PER_BLOCK = 3
|
||||
_SPATIAL_ALIGNMENT = 16
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LingBotPaths:
|
||||
image: Path
|
||||
action_dir: Path
|
||||
action_root: Path
|
||||
action_relative: Path
|
||||
camera_frames: int
|
||||
output: Path
|
||||
|
||||
|
||||
def _validate_camera_arrays(action_dir: Path) -> int:
|
||||
"""Validate the bounded camera contract before starting the engine."""
|
||||
|
||||
try:
|
||||
poses = np.load(action_dir / "poses.npy", mmap_mode="r", allow_pickle=False)
|
||||
intrinsics = np.load(action_dir / "intrinsics.npy", mmap_mode="r", allow_pickle=False)
|
||||
except (OSError, TypeError, ValueError):
|
||||
raise ValueError("Camera files must be valid numeric poses.npy and intrinsics.npy arrays.") from None
|
||||
if poses.ndim != 3 or poses.shape[1:] != (4, 4):
|
||||
raise ValueError("poses.npy must have shape [frames, 4, 4].")
|
||||
if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,):
|
||||
raise ValueError("intrinsics.npy must have shape [frames, 4].")
|
||||
if poses.shape[0] != intrinsics.shape[0]:
|
||||
raise ValueError("poses.npy and intrinsics.npy must contain the same frame count.")
|
||||
if not 1 <= poses.shape[0] <= _MAX_SOURCE_CAMERA_FRAMES:
|
||||
raise ValueError("Camera arrays must contain between 1 and 4096 source frames.")
|
||||
if poses.dtype.kind not in "fiu" or intrinsics.dtype.kind not in "fiu":
|
||||
raise ValueError("Camera arrays must contain real numeric values.")
|
||||
if not np.isfinite(poses).all() or not np.isfinite(intrinsics).all():
|
||||
raise ValueError("Camera arrays must contain only finite values.")
|
||||
return int(poses.shape[0])
|
||||
|
||||
|
||||
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate a LingBot-World v2 video from a first-frame image and camera trajectory."
|
||||
)
|
||||
parser.add_argument("--model", default=_MODEL, help="Official Hugging Face model ID or local checkpoint path.")
|
||||
parser.add_argument("--prompt", required=True, help="Text description of the scene and desired motion.")
|
||||
parser.add_argument("--image", required=True, help="Path to the first-frame image.")
|
||||
parser.add_argument(
|
||||
"--action-dir",
|
||||
required=True,
|
||||
help="Camera action directory containing poses.npy and intrinsics.npy.",
|
||||
)
|
||||
parser.add_argument("--height", type=int, default=480, help="Output height; must be divisible by 16.")
|
||||
parser.add_argument("--width", type=int, default=832, help="Output width; must be divisible by 16.")
|
||||
parser.add_argument(
|
||||
"--num-frames",
|
||||
type=int,
|
||||
default=9,
|
||||
help="Raw output frame count. Nine frames are one three-latent-frame causal block.",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument(
|
||||
"--tensor-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of GPUs used for tensor parallelism inside the DiT.",
|
||||
)
|
||||
parser.add_argument("--flow-shift", type=float, default=5.0, help="Positive FlowUniPC scheduler shift.")
|
||||
parser.add_argument("--fps", type=int, default=16, help="Frames per second in the exported MP4.")
|
||||
parser.add_argument("--output", default="lingbot_world_v2.mp4", help="Output MP4 path.")
|
||||
parser.add_argument(
|
||||
"--enforce-eager",
|
||||
action="store_true",
|
||||
help="Disable compilation and force eager execution.",
|
||||
)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def resolve_cli_paths(args: argparse.Namespace) -> LingBotPaths:
|
||||
"""Canonicalize local inputs and construct a safe trusted-root mapping."""
|
||||
|
||||
image = Path(args.image).expanduser().resolve()
|
||||
if not image.is_file():
|
||||
raise ValueError("--image must point to an existing file.")
|
||||
|
||||
action_dir = Path(args.action_dir).expanduser().resolve()
|
||||
if not action_dir.is_dir():
|
||||
raise ValueError("--action-dir must point to an existing directory.")
|
||||
if not all((action_dir / name).is_file() for name in ("poses.npy", "intrinsics.npy")):
|
||||
raise ValueError("--action-dir must contain poses.npy and intrinsics.npy.")
|
||||
camera_frames = _validate_camera_arrays(action_dir)
|
||||
|
||||
action_root = action_dir.parent.resolve()
|
||||
try:
|
||||
action_relative = action_dir.relative_to(action_root)
|
||||
except ValueError:
|
||||
raise ValueError("--action-dir must be contained by its canonical trusted action root.") from None
|
||||
if not action_relative.parts:
|
||||
raise ValueError("--action-dir must name a directory contained by a trusted action root.")
|
||||
|
||||
output = Path(args.output).expanduser().resolve()
|
||||
return LingBotPaths(
|
||||
image=image,
|
||||
action_dir=action_dir,
|
||||
action_root=action_root,
|
||||
action_relative=action_relative,
|
||||
camera_frames=camera_frames,
|
||||
output=output,
|
||||
)
|
||||
|
||||
|
||||
def _positive_finite(value: float, flag: str) -> float:
|
||||
result = float(value)
|
||||
if not math.isfinite(result) or result <= 0:
|
||||
raise ValueError(f"{flag} must be a positive finite number.")
|
||||
return result
|
||||
|
||||
|
||||
def build_omni_kwargs(
|
||||
args: argparse.Namespace,
|
||||
paths: LingBotPaths,
|
||||
*,
|
||||
parallel_config: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""Build Omni configuration without overriding checkpoint class discovery."""
|
||||
|
||||
if args.tensor_parallel_size <= 0:
|
||||
raise ValueError("--tensor-parallel-size must be a positive integer.")
|
||||
flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
|
||||
model_path = Path(args.model).expanduser()
|
||||
model = str(model_path.resolve()) if model_path.exists() else args.model
|
||||
return {
|
||||
"model": model,
|
||||
"flow_shift": flow_shift,
|
||||
"parallel_config": parallel_config,
|
||||
"enforce_eager": args.enforce_eager,
|
||||
"model_config": {"lingbot_action_root": str(paths.action_root)},
|
||||
}
|
||||
|
||||
|
||||
def build_request(
|
||||
args: argparse.Namespace,
|
||||
paths: LingBotPaths,
|
||||
) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
"""Build the prompt and fixed v1 causal-DMD sampling contract."""
|
||||
|
||||
prompt_text = args.prompt.strip()
|
||||
if not prompt_text:
|
||||
raise ValueError("--prompt must contain non-whitespace text.")
|
||||
if args.height <= 0 or args.width <= 0:
|
||||
raise ValueError("--height and --width must be positive integers.")
|
||||
if args.height % _SPATIAL_ALIGNMENT or args.width % _SPATIAL_ALIGNMENT:
|
||||
raise ValueError("--height and --width must each be divisible by 16.")
|
||||
if args.height * args.width > _MAX_PIXEL_AREA:
|
||||
raise ValueError("--height * --width must not exceed 480 * 832 pixels.")
|
||||
if args.num_frames <= 0:
|
||||
raise ValueError("--num-frames must be a positive integer.")
|
||||
if args.num_frames > _MAX_RAW_FRAMES:
|
||||
raise ValueError("--num-frames must not exceed 117 raw frames.")
|
||||
if args.num_frames > paths.camera_frames:
|
||||
raise ValueError("Camera arrays must contain at least --num-frames frames.")
|
||||
if (args.num_frames - 1) % _TEMPORAL_COMPRESSION:
|
||||
raise ValueError("--num-frames must satisfy (num_frames - 1) divisible by 4.")
|
||||
latent_frames = (args.num_frames - 1) // _TEMPORAL_COMPRESSION + 1
|
||||
if latent_frames % _LATENT_FRAMES_PER_BLOCK:
|
||||
raise ValueError("--num-frames must map to whole three-frame latent blocks.")
|
||||
if args.fps <= 0:
|
||||
raise ValueError("--fps must be a positive integer.")
|
||||
flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
|
||||
|
||||
prompt = {
|
||||
"prompt": prompt_text,
|
||||
"multi_modal_data": {"image": str(paths.image)},
|
||||
}
|
||||
sampling_kwargs = {
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_frames": args.num_frames,
|
||||
"num_inference_steps": _NUM_INFERENCE_STEPS,
|
||||
"max_sequence_length": _MAX_SEQUENCE_LENGTH,
|
||||
"seed": args.seed,
|
||||
"fps": args.fps,
|
||||
"extra_args": {
|
||||
"action_path": paths.action_relative.as_posix(),
|
||||
"flow_shift": flow_shift,
|
||||
},
|
||||
}
|
||||
return prompt, sampling_kwargs
|
||||
|
||||
|
||||
def extract_video_array(outputs: Any) -> np.ndarray:
|
||||
"""Unwrap one standard Omni diffusion video into ``[frames, H, W, C]``."""
|
||||
|
||||
if not isinstance(outputs, list) or len(outputs) != 1:
|
||||
raise ValueError("Expected exactly one Omni output for the single LingBot request.")
|
||||
output = outputs[0]
|
||||
request_output = getattr(output, "request_output", None)
|
||||
if request_output is not None and getattr(request_output, "images", None):
|
||||
output = request_output
|
||||
|
||||
images = getattr(output, "images", None)
|
||||
if not images:
|
||||
multimodal_output = getattr(output, "multimodal_output", None) or {}
|
||||
images = multimodal_output.get("video")
|
||||
if images is None:
|
||||
raise ValueError("No video frames found in Omni output.")
|
||||
if isinstance(images, list) and len(images) == 1:
|
||||
images = images[0]
|
||||
|
||||
video: np.ndarray = np.asarray(images)
|
||||
if video.ndim == 5:
|
||||
if video.shape[0] != 1:
|
||||
raise ValueError(f"Expected one generated video, got batch shape {video.shape}.")
|
||||
video = video[0]
|
||||
if video.ndim != 4 or video.shape[-1] not in (3, 4):
|
||||
raise ValueError(f"Expected video shape [frames, height, width, channels], got {video.shape}.")
|
||||
return video
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> Path:
|
||||
args = parse_args(argv)
|
||||
paths = resolve_cli_paths(args)
|
||||
prompt, sampling_kwargs = build_request(args, paths)
|
||||
paths.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Keep vLLM imports below pure CLI/path validation so ``--help`` and the
|
||||
# helper tests work on machines without the native vLLM runtime.
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
parallel_config = DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size)
|
||||
omni_kwargs = build_omni_kwargs(args, paths, parallel_config=parallel_config)
|
||||
omni = Omni(**omni_kwargs)
|
||||
try:
|
||||
outputs = omni.generate(
|
||||
prompt,
|
||||
OmniDiffusionSamplingParams(**sampling_kwargs),
|
||||
use_tqdm=False,
|
||||
)
|
||||
finally:
|
||||
omni.close()
|
||||
|
||||
frames = extract_video_array(outputs)
|
||||
export_to_video(frames, str(paths.output), fps=args.fps)
|
||||
print(f"Saved {frames.shape[0]} frames to {paths.output}")
|
||||
return paths.output
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,239 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Run LingBot-World 2.0 as an in-process realtime AR-Diffusion session.
|
||||
|
||||
This example exercises the internal session API. It is not a public HTTP or
|
||||
WebSocket protocol. Each JSONL event advances the world by one three-latent
|
||||
frame AR block and writes the latent plus its identity metadata. The current
|
||||
scoped implementation supports at most ten ticks per generation epoch because
|
||||
its image condition is bounded to 117 pixel frames; reset or create a session
|
||||
to start a new world.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import math
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
|
||||
_CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1"
|
||||
_FRAMES_PER_BLOCK = 3
|
||||
# ((117 pixel frames - 1) / VAE temporal factor 4 + 1) / 3 latent frames.
|
||||
_MAX_REALTIME_TICKS = 10
|
||||
|
||||
|
||||
def _camera_event_data(frames: list[list[str]]) -> dict[str, Any]:
|
||||
"""Build the event-side script consumed by LingBotCameraControlReducer."""
|
||||
return {"mode": "script", "frames": frames}
|
||||
|
||||
|
||||
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Run an in-process realtime LingBot-World 2.0 session.")
|
||||
parser.add_argument("--model", default=_MODEL, help="Hugging Face model ID or local checkpoint path.")
|
||||
parser.add_argument("--image", required=True, help="Initial RGB image.")
|
||||
parser.add_argument("--prompt", required=True, help="Initial scene prompt.")
|
||||
parser.add_argument(
|
||||
"--events",
|
||||
required=True,
|
||||
help="JSONL file with one prompt/action event per AR block; at most 10 events.",
|
||||
)
|
||||
parser.add_argument("--output-dir", required=True, help="Directory for chunk latents and metadata.")
|
||||
parser.add_argument("--session-id", default="lingbot-world", help="Persistent world session identifier.")
|
||||
parser.add_argument("--height", type=int, default=480)
|
||||
parser.add_argument("--width", type=int, default=832)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--gpu-memory-fraction", type=float, default=0.1)
|
||||
parser.add_argument("--tensor-parallel-size", type=int, default=1)
|
||||
parser.add_argument("--enforce-eager", action="store_true")
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def _load_events(path: Path) -> list[dict[str, Any]]:
|
||||
events: list[dict[str, Any]] = []
|
||||
for line_number, raw_line in enumerate(path.read_text().splitlines(), start=1):
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
value = json.loads(line)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"events line {line_number} is not valid JSON: {exc.msg}.") from None
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError(f"events line {line_number} must be a JSON object.")
|
||||
event_id = value.get("event_id")
|
||||
if isinstance(event_id, bool) or not isinstance(event_id, int) or event_id < 0:
|
||||
raise ValueError(f"events line {line_number} requires a non-negative integer event_id.")
|
||||
prompt = value.get("prompt")
|
||||
if prompt is not None and (not isinstance(prompt, str) or not prompt.strip()):
|
||||
raise ValueError(f"events line {line_number} prompt must be a non-empty string.")
|
||||
frames = value.get("frames")
|
||||
if frames is not None:
|
||||
if not isinstance(frames, list) or len(frames) != _FRAMES_PER_BLOCK:
|
||||
raise ValueError(f"events line {line_number} frames must contain exactly three lists.")
|
||||
for frame in frames:
|
||||
if not isinstance(frame, list) or any(not isinstance(action, str) for action in frame):
|
||||
raise ValueError(f"events line {line_number} frames must contain only action strings.")
|
||||
if prompt is None and frames is None:
|
||||
raise ValueError(f"events line {line_number} must update prompt and/or frames.")
|
||||
events.append({"event_id": event_id, "prompt": prompt, "frames": frames})
|
||||
if not events:
|
||||
raise ValueError("events file must contain at least one event.")
|
||||
if len(events) > _MAX_REALTIME_TICKS:
|
||||
raise ValueError(
|
||||
"events file must contain at most "
|
||||
f"{_MAX_REALTIME_TICKS} events because the current LingBot realtime "
|
||||
"image-condition horizon is 117 pixel frames."
|
||||
)
|
||||
event_ids = [event["event_id"] for event in events]
|
||||
if event_ids != sorted(set(event_ids)):
|
||||
raise ValueError("event_id values must be unique and strictly increasing.")
|
||||
return events
|
||||
|
||||
|
||||
def _validate_args(args: argparse.Namespace) -> tuple[Path, Path, Path]:
|
||||
image = Path(args.image).expanduser().resolve()
|
||||
events = Path(args.events).expanduser().resolve()
|
||||
output_dir = Path(args.output_dir).expanduser().resolve()
|
||||
if not image.is_file():
|
||||
raise ValueError("--image must point to an existing file.")
|
||||
if not events.is_file():
|
||||
raise ValueError("--events must point to an existing JSONL file.")
|
||||
if not args.prompt.strip():
|
||||
raise ValueError("--prompt must contain non-whitespace text.")
|
||||
if args.height <= 0 or args.width <= 0 or args.height % 16 or args.width % 16:
|
||||
raise ValueError("--height and --width must be positive multiples of 16.")
|
||||
if args.tensor_parallel_size <= 0:
|
||||
raise ValueError("--tensor-parallel-size must be positive.")
|
||||
if not math.isfinite(args.gpu_memory_fraction) or not 0 < args.gpu_memory_fraction <= 1:
|
||||
raise ValueError("--gpu-memory-fraction must be in (0, 1].")
|
||||
return image, events, output_dir
|
||||
|
||||
|
||||
async def run(argv: Sequence[str] | None = None) -> Path:
|
||||
args = parse_args(argv)
|
||||
image, events_path, output_dir = _validate_args(args)
|
||||
events = _load_events(events_path)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Imports stay below pure input validation so --help and helper tests do not
|
||||
# require a CUDA-enabled vLLM installation.
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.models.lingbot_world.actions import LingBotCameraControlReducer
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
|
||||
from vllm_omni.experimental.ar_diffusion.session import (
|
||||
ARDiffusionSessionEvent,
|
||||
ARDiffusionSessionManager,
|
||||
ARDiffusionWorkerLifecycle,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionControlInput
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
engine = AsyncOmni(
|
||||
model=args.model,
|
||||
engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine",
|
||||
enforce_eager=args.enforce_eager,
|
||||
parallel_config=DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size),
|
||||
max_num_seqs=1,
|
||||
model_config={
|
||||
"ar_diffusion_height": args.height,
|
||||
"ar_diffusion_width": args.width,
|
||||
"ar_diffusion_kv_config": {
|
||||
"gpu_memory_fraction": args.gpu_memory_fraction,
|
||||
"warmup_cudagraph": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
sampling = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=9,
|
||||
num_inference_steps=4,
|
||||
max_sequence_length=512,
|
||||
seed=args.seed,
|
||||
output_type="latent",
|
||||
extra_args={"flow_shift": 5.0},
|
||||
)
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
engine,
|
||||
prompt_provider=lambda tick: {
|
||||
"prompt": tick.prompt,
|
||||
"multi_modal_data": {"image": str(image)},
|
||||
},
|
||||
sampling_params_list=[sampling],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
manager = ARDiffusionSessionManager(
|
||||
tick_consumer=consumer,
|
||||
lifecycle=ARDiffusionWorkerLifecycle(engine, stage_ids=[0], timeout=180.0),
|
||||
max_pending_events=32,
|
||||
control_reducer_factory=LingBotCameraControlReducer,
|
||||
)
|
||||
session = await manager.create_session(args.session_id)
|
||||
measurements: list[dict[str, Any]] = []
|
||||
try:
|
||||
for chunk_index, event in enumerate(events):
|
||||
controls = ()
|
||||
if event["frames"] is not None:
|
||||
controls = (
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema=_CAMERA_ACTION_SCHEMA,
|
||||
data=_camera_event_data(event["frames"]),
|
||||
),
|
||||
)
|
||||
await session.accept_event(
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=event["event_id"],
|
||||
prompt=event["prompt"]
|
||||
if event["prompt"] is not None
|
||||
else (args.prompt if chunk_index == 0 else None),
|
||||
controls=controls,
|
||||
)
|
||||
)
|
||||
started = time.perf_counter()
|
||||
output = await session.next_chunk()
|
||||
elapsed = time.perf_counter() - started
|
||||
if len(output.images) != 1 or not isinstance(output.images[0], torch.Tensor):
|
||||
raise RuntimeError("Expected one latent tensor from each realtime LingBot chunk.")
|
||||
latent = output.images[0].detach().float().cpu()
|
||||
metadata = consumer.chunk_metadata(output).to_dict()
|
||||
torch.save(latent, output_dir / f"chunk_{chunk_index:03d}.pt")
|
||||
(output_dir / f"chunk_{chunk_index:03d}.json").write_text(
|
||||
json.dumps(metadata, indent=2, sort_keys=True) + "\n"
|
||||
)
|
||||
measurements.append(
|
||||
{
|
||||
"chunk_index": chunk_index,
|
||||
"latency_seconds": elapsed,
|
||||
"shape": list(latent.shape),
|
||||
"finite": bool(torch.isfinite(latent).all()),
|
||||
"metadata": metadata,
|
||||
}
|
||||
)
|
||||
print(json.dumps(measurements[-1], sort_keys=True), flush=True)
|
||||
finally:
|
||||
try:
|
||||
await manager.close_session(args.session_id)
|
||||
finally:
|
||||
engine.shutdown()
|
||||
|
||||
summary_path = output_dir / "summary.json"
|
||||
summary_path.write_text(json.dumps({"chunks": measurements}, indent=2, sort_keys=True) + "\n")
|
||||
return summary_path
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> Path:
|
||||
return asyncio.run(run(argv))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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}")
|
||||
|
||||
@@ -96,7 +96,7 @@ if __name__ == "__main__":
|
||||
),
|
||||
)
|
||||
from diffusers.utils import export_to_video
|
||||
frames = outputs[0].request_output.images
|
||||
frames = outputs[0].images
|
||||
export_to_video(frames, "quick_test_i2v.mp4", fps=16)
|
||||
```
|
||||
|
||||
|
||||
@@ -68,6 +68,7 @@ from vllm_omni.model_extras import (
|
||||
from vllm_omni.model_extras import (
|
||||
get_extra_body_params,
|
||||
get_model_class_name,
|
||||
should_preserve_reference_image_size,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
@@ -454,7 +455,13 @@ def main():
|
||||
|
||||
media_inputs: dict[str, Any] = {}
|
||||
if image is not None:
|
||||
media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
preserve_image_size = should_preserve_reference_image_size(
|
||||
model_class_name,
|
||||
model=args.model,
|
||||
)
|
||||
media_inputs["image"] = (
|
||||
image if preserve_image_size else image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
)
|
||||
if last_image is not None:
|
||||
media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
if mask_image is not None:
|
||||
|
||||
@@ -82,9 +82,6 @@ def parse_args():
|
||||
# ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
|
||||
# and ``create_default_diffusion`` materializes the stage config.
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(p)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -335,7 +335,7 @@ def main(args):
|
||||
print(f"use_tts (tts_bos): {use_tts}")
|
||||
|
||||
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
|
||||
@@ -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,
|
||||
|
||||
+60
-25
@@ -1,21 +1,32 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Reproduce MiniMax-H3 request-mode and DLO/DP2 lifecycle behavior.
|
||||
"""Validate MiniMax-H3 request-mode and DLO lifecycle behavior.
|
||||
|
||||
The DLO modes accept any positive data-parallel size. With AllGather enabled,
|
||||
the script submits one complete DP wave and verifies recovery after an invalid
|
||||
wave. It also checks that every mode shuts down all worker processes cleanly.
|
||||
|
||||
Examples:
|
||||
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
|
||||
VLLM_OMNI_DLO_DP_WAVE_TIMEOUT=600 \
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode dlo --dp-size 4
|
||||
|
||||
# Rank-local mmap storage, with no DLO collective or DP wave scheduling.
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
|
||||
CUDA_VISIBLE_DEVICES=0,1 \
|
||||
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
|
||||
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode dlo-no-allgather --dp-size 2
|
||||
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
|
||||
CUDA_VISIBLE_DEVICES=0,1 \
|
||||
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode request
|
||||
python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
|
||||
--model /path/to/MiniMax-H3/FL2VA --mode request --tp-size 2
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -40,14 +51,33 @@ DEFAULT_PROMPTS = (
|
||||
)
|
||||
|
||||
|
||||
def positive_int(value: str) -> int:
|
||||
parsed = int(value)
|
||||
if parsed < 1:
|
||||
raise argparse.ArgumentTypeError("must be a positive integer")
|
||||
return parsed
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA")
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=("dlo-dp2", "request"),
|
||||
choices=("dlo", "dlo-no-allgather", "request"),
|
||||
required=True,
|
||||
help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode",
|
||||
help="DLO (with or without AllGather) or ordinary non-DLO request mode",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dp-size",
|
||||
type=positive_int,
|
||||
default=2,
|
||||
help="DLO data-parallel size (for example 1, 2, 4, or 8; default: 2)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tp-size",
|
||||
type=positive_int,
|
||||
default=2,
|
||||
help="Tensor-parallel size used only by request mode (default: 2)",
|
||||
)
|
||||
parser.add_argument("--steps", type=int, default=2)
|
||||
parser.add_argument("--duration", type=float, default=5.0)
|
||||
@@ -60,10 +90,11 @@ def parse_args() -> argparse.Namespace:
|
||||
|
||||
|
||||
def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
|
||||
is_dlo = args.mode in {"dlo", "dlo-no-allgather"}
|
||||
common: dict[str, Any] = {
|
||||
"model": args.model,
|
||||
"trust_remote_code": True,
|
||||
"num_gpus": 2,
|
||||
"num_gpus": args.dp_size if is_dlo else args.tp_size,
|
||||
"usp": 1,
|
||||
"ring": 1,
|
||||
"vae_parallel_mode": "tile",
|
||||
@@ -74,22 +105,22 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
|
||||
"stage_init_timeout": args.init_timeout,
|
||||
"init_timeout": args.init_timeout,
|
||||
}
|
||||
if args.mode == "dlo-dp2":
|
||||
if is_dlo:
|
||||
common.update(
|
||||
tensor_parallel_size=1,
|
||||
data_parallel_size=2,
|
||||
data_parallel_size=args.dp_size,
|
||||
text_encoder_tp_size=1,
|
||||
vae_patch_parallel_size=1,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
dlo_use_allgather=args.mode == "dlo",
|
||||
dlo_resident_layers=0,
|
||||
)
|
||||
else:
|
||||
common.update(
|
||||
tensor_parallel_size=2,
|
||||
tensor_parallel_size=args.tp_size,
|
||||
data_parallel_size=1,
|
||||
text_encoder_tp_size=2,
|
||||
vae_patch_parallel_size=2,
|
||||
text_encoder_tp_size=args.tp_size,
|
||||
vae_patch_parallel_size=args.tp_size,
|
||||
enable_distributed_layerwise_offload=False,
|
||||
)
|
||||
return common
|
||||
@@ -172,7 +203,8 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
|
||||
"engine_kwargs": engine_kwargs(args),
|
||||
}
|
||||
try:
|
||||
if args.mode == "dlo-dp2":
|
||||
uses_dp_wave = args.mode == "dlo" and args.dp_size > 1
|
||||
if uses_dp_wave:
|
||||
started = time.perf_counter()
|
||||
asymmetric = await asyncio.gather(
|
||||
generate_one(
|
||||
@@ -182,12 +214,15 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
|
||||
prompt="",
|
||||
seed=1001,
|
||||
),
|
||||
generate_one(
|
||||
engine,
|
||||
args,
|
||||
request_id="invalid-peer",
|
||||
prompt=DEFAULT_PROMPTS[0],
|
||||
seed=1002,
|
||||
*(
|
||||
generate_one(
|
||||
engine,
|
||||
args,
|
||||
request_id=f"invalid-peer-{index}",
|
||||
prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
|
||||
seed=1001 + index,
|
||||
)
|
||||
for index in range(1, args.dp_size)
|
||||
),
|
||||
return_exceptions=True,
|
||||
)
|
||||
@@ -195,18 +230,18 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
|
||||
summary["asymmetric_errors"] = [
|
||||
f"{type(result).__name__}: {result}" for result in asymmetric if isinstance(result, BaseException)
|
||||
]
|
||||
if len(summary["asymmetric_errors"]) != 2:
|
||||
raise RuntimeError("Both requests in the asymmetric DP wave must fail before dispatch")
|
||||
if len(summary["asymmetric_errors"]) != args.dp_size:
|
||||
raise RuntimeError("Every request in the asymmetric DP wave must fail before dispatch")
|
||||
|
||||
started = time.perf_counter()
|
||||
request_count = 2 if args.mode == "dlo-dp2" else 1
|
||||
request_count = args.dp_size if uses_dp_wave else 1
|
||||
outputs = await asyncio.gather(
|
||||
*(
|
||||
generate_one(
|
||||
engine,
|
||||
args,
|
||||
request_id=f"recovery-{index}",
|
||||
prompt=DEFAULT_PROMPTS[index],
|
||||
prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
|
||||
seed=2000 + index,
|
||||
)
|
||||
for index in range(request_count)
|
||||
@@ -67,7 +67,7 @@ def main() -> None:
|
||||
parser.add_argument("--output-dir", default="results/nemotron_voicechat")
|
||||
parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
"--deploy-config",
|
||||
default=None,
|
||||
help=(
|
||||
"Deploy yaml override. Pass vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml "
|
||||
@@ -134,8 +134,8 @@ def main() -> None:
|
||||
]
|
||||
|
||||
omni_kwargs: dict = {}
|
||||
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
|
||||
omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs)
|
||||
outputs = omni.generate(inputs, sampling_params_list)
|
||||
|
||||
|
||||
@@ -484,10 +484,10 @@ def parse_args():
|
||||
help="Path to a .txt file with one prompt per line (preferred).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a stage configs file.",
|
||||
help="Path to a deploy config YAML.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video-path",
|
||||
|
||||
@@ -232,15 +232,8 @@ def main(args):
|
||||
raise ValueError(f"Unknown query type: {args.query_type}")
|
||||
|
||||
# Initialize vLLM-Omni with Step-Audio2
|
||||
# Resolve deploy config path. Keep explicit stage_configs_path for legacy
|
||||
# custom configs, but use bundled deploy configs by default.
|
||||
deploy_config_path = None
|
||||
stage_config_path = None
|
||||
if args.stage_configs_path and args.deploy_config:
|
||||
raise ValueError("--stage-configs-path and --deploy-config are mutually exclusive")
|
||||
if args.stage_configs_path:
|
||||
stage_config_path = args.stage_configs_path
|
||||
elif args.deploy_config:
|
||||
# Resolve an explicit deploy config or use the bundled task default.
|
||||
if args.deploy_config:
|
||||
deploy_config_path = args.deploy_config
|
||||
else:
|
||||
configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/deploy"
|
||||
@@ -262,10 +255,7 @@ def main(args):
|
||||
"ray_address": args.ray_address,
|
||||
"trust_remote_code": True,
|
||||
}
|
||||
if stage_config_path is not None:
|
||||
omni_kwargs["stage_configs_path"] = stage_config_path
|
||||
else:
|
||||
omni_kwargs["deploy_config"] = deploy_config_path
|
||||
omni_kwargs["deploy_config"] = deploy_config_path
|
||||
|
||||
omni_llm = Omni(**omni_kwargs)
|
||||
|
||||
@@ -418,12 +408,6 @@ def parse_args():
|
||||
" - HuggingFace ID: stepfun-ai/Step-Audio-2-mini"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a legacy stage config YAML file (default: use bundled deploy config)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
|
||||
@@ -74,7 +74,7 @@ if __name__ == "__main__":
|
||||
omni = Omni(model="Qwen/Qwen-Image")
|
||||
prompt = "a cup of coffee on the table"
|
||||
outputs = omni.generate(prompt)
|
||||
images = outputs[0].request_output.images
|
||||
images = outputs[0].images
|
||||
images[0].save("coffee.png")
|
||||
```
|
||||
|
||||
@@ -265,7 +265,7 @@ if __name__ == "__main__":
|
||||
]
|
||||
outputs = omni.generate(prompts)
|
||||
for i, output in enumerate(outputs):
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
!!! info
|
||||
@@ -301,7 +301,7 @@ if __name__ == "__main__":
|
||||
}
|
||||
])
|
||||
for i, output in enumerate(outputs):
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
You can also pass a negative prompt via the CLI argument `--negative-prompt`:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -16,6 +16,7 @@ list of supported architectures across all modalities, see
|
||||
|---|---|---|---|---|---|---|
|
||||
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
|
||||
| Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning WIP) | — (serving WIP) | zero-shot | 22.05 kHz |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
|
||||
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz |
|
||||
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
|
||||
@@ -23,6 +24,7 @@ list of supported architectures across all modalities, see
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz |
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz |
|
||||
| dots.tts | `rednote-hilab/dots.tts-soar` | single (native AR) | — (not wired yet) | — | — | 48 kHz |
|
||||
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz |
|
||||
| IndexTTS-2.5 | native `checkpoints/` bundle | 2 (AR talker + EnhancedCodec + S2Mel DiT + BigVGAN) | ✓ (required) | — | multilingual (`--lang`) + emotion control | 22.05 kHz |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
|
||||
@@ -394,6 +396,58 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
|
||||
|
||||
---
|
||||
|
||||
## Gepard-1.0
|
||||
|
||||
Single-stage native AR TTS at 22.05 kHz. Pipeline: `Qwen3.5 backbone → 32 FSQ codebook heads (one frame/step) → NeMo NanoCodec`. Backbone runs under vLLM paged attention; the 32-head sampling + learned embedding feedback ride the native-AR runner hooks. `enforce_eager` (CUDA graph is a perf follow-up).
|
||||
|
||||
### Prerequisites
|
||||
The NanoCodec decoder needs NeMo (NVIDIA Open Model License), installed separately:
|
||||
|
||||
```bash
|
||||
uv pip install "nemo_toolkit[tts]==2.7.3" \
|
||||
--overrides <(printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0")
|
||||
```
|
||||
|
||||
The overrides are what make that command work. `nemo_toolkit[tts]` declares
|
||||
`transformers~=4.57.0`, so a plain install downgrades the version
|
||||
`requirements/common.txt` pins and takes `huggingface-hub` back to 0.x with it — after
|
||||
which `import vllm` fails. Both packages are named here rather than left to resolve on
|
||||
their own: transformers 5.x happens to require `huggingface-hub` 1.x, but relying on
|
||||
that would make the second repair an accident of the first. `--constraint` cannot be
|
||||
used instead, since it narrows a version range rather than overriding a dependency a
|
||||
package declares, and the two sets never intersect. NeMo 2.7.3 decodes correctly
|
||||
against transformers 5.x, so nothing is lost by holding it there.
|
||||
|
||||
(The same clash is why this is documented here rather than declared in
|
||||
`pyproject.toml`: an extra resolves together with the base dependencies, so
|
||||
`vllm-omni[gepard]` would be unsatisfiable rather than opt-in. SoulX-Singer English
|
||||
SVS documents `nemo_toolkit[asr]` the same way.)
|
||||
|
||||
On a host whose CUDA toolkit cannot build kernels for the GPU — no `nvcc`/`ninja`, or a
|
||||
consumer Blackwell (`sm_120`) card — also set `VLLM_USE_FLASHINFER_SAMPLER=0`. vLLM's
|
||||
sampler JIT-compiles through FlashInfer on first use, and where that build fails the engine
|
||||
dies inside `profile_run` long after every import has succeeded, which reads as a model
|
||||
failure rather than an environment one. The PyTorch sampler it falls back to costs nothing
|
||||
here.
|
||||
|
||||
### Quick start (zero-shot, default voice)
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/gepard/end2end.py \
|
||||
--text "Hello, this is a Gepard demo."
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
Not yet — PR1 is zero-shot only (the learned `null_prefix` default voice). Reference-audio cloning (`ref_compressor` speaker prefix) is a follow-up PR.
|
||||
|
||||
### Notes
|
||||
- Output: 22.05 kHz mono WAV.
|
||||
- First run downloads two checkpoints, not one: the model itself, and the NanoCodec decoder it names in `codec_id` (`nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps` by default). Both are fetched automatically; expect the first startup to be correspondingly slower.
|
||||
- Offline only for now; the `/v1/audio/speech` adapter is a follow-up PR.
|
||||
- Deploy config: `vllm_omni/deploy/gepard.yaml` (the example's default; copy it and pass `--deploy-config` to change it).
|
||||
- Generation length and reproducibility are stage settings, not script flags. One output token is one audio frame, so `max_tokens` in the YAML is the frame budget; `seed` makes the in-model 32-head sampling reproducible. The script deliberately passes no `SamplingParams`: one supplied by a caller replaces the stage defaults wholesale rather than merging, which would drop the pipeline's stop token and run every request to `max_tokens`.
|
||||
- Text length: short texts are repeated internally to match the training layout (the checkpoint's `text_repetition` block); the upper bound is the stage's `max_model_len`, enforced by the engine. Empty text is rejected rather than voiced.
|
||||
- `VLLM_GEPARD_GREEDY=1` swaps the 32-head Gumbel-max sampling for argmax, for reproducible comparisons that do not depend on a seed.
|
||||
|
||||
## VoxCPM2
|
||||
|
||||
Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`.
|
||||
@@ -430,6 +484,31 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [
|
||||
|
||||
---
|
||||
|
||||
## dots.tts
|
||||
|
||||
Single-stage native AR TTS at 48 kHz (rednote-hilab). Pipeline: `Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR loopback → AudioVAE (streaming decode)`. Same "vLLM-native base LM + side-path computation" pattern as VoxCPM2, with a plain Qwen2 backbone instead of MiniCPM4 and no FSQ / residual LM stage.
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
|
||||
--model rednote-hilab/dots.tts-soar \
|
||||
--text "Hello, this is a test of dots TTS running on vLLM Omni."
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
Not wired in this release — generation is zero-shot only. The CAM++ x-vector speaker encoder weights load, but `end2end.py` has no `--ref-audio`/`--ref-text` flags yet.
|
||||
|
||||
### Streaming
|
||||
The AudioVAE decoder has an internal streaming path (`init_stream_state` / `stream_step` / `stream_flush`) used to avoid boundary artifacts between 160 ms patches, but it is not yet exposed through an online serving endpoint or example.
|
||||
|
||||
### Notes
|
||||
- Output: 48 kHz mono WAV.
|
||||
- Deploy config: `vllm_omni/deploy/dots_tts.yaml` (auto-loaded by HF `model_type`).
|
||||
- Checkpoints: `rednote-hilab/dots.tts-soar` is the validated default. `dots.tts-base` shares the same architecture but is unvalidated in this repo. `dots.tts-mf` (MeanFlow, 2-4 step) is not supported yet.
|
||||
- Known limitation: no CUDA graph capture and no batched side-path computation yet, so concurrent requests do not currently scale (each request's DiT Euler steps run serially). See `recipes/rednote-hilab/dots.tts.md` for details and the roadmap.
|
||||
|
||||
---
|
||||
|
||||
## IndexTTS-2 and IndexTTS-2.5
|
||||
|
||||
Both versions use a 2-stage TTS pipeline and produce 22.05 kHz mono audio. Stage 0 is a GPT AR talker. IndexTTS-2 Stage 1 uses RepCodec semantic embeddings, GPT latent, S2Mel CFM/DiT, and BigVGAN. IndexTTS-2.5 replaces RepCodec with EnhancedCodec and uses the official code-only path (`use_gpt_latent=false`). Every request requires reference audio for zero-shot voice cloning.
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
"""Offline dots.tts inference example (single-stage native AR pipeline).
|
||||
|
||||
Uses the single-stage native AR config (vllm_omni/deploy/dots_tts.yaml).
|
||||
Reference checkpoint: rednote-hilab/dots.tts-soar.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
SAMPLE_RATE = 48_000
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(description="Offline dots.tts native AR inference")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="rednote-hilab/dots.tts-soar",
|
||||
help="dots.tts model path or HuggingFace repo ID.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
type=str,
|
||||
default="Hello, this is a test of dots TTS running on vLLM Omni.",
|
||||
help="Text to synthesize.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default="output_audio",
|
||||
help="Directory for output WAV files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Override the deploy config path. If unset, auto-loads "
|
||||
"vllm_omni/deploy/dots_tts.yaml based on the HF model_type.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def extract_audio(multimodal_output: dict) -> torch.Tensor:
|
||||
"""Extract the final complete audio tensor from multimodal output.
|
||||
|
||||
The output processor concatenates per-step delta tensors under
|
||||
``model_outputs``. Falls back to ``audio`` for backwards compat.
|
||||
"""
|
||||
audio = multimodal_output.get("model_outputs")
|
||||
if audio is None:
|
||||
audio = multimodal_output.get("audio")
|
||||
if audio is None:
|
||||
raise ValueError(f"No audio key in multimodal_output: {list(multimodal_output.keys())}")
|
||||
|
||||
if isinstance(audio, list):
|
||||
valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
|
||||
if not valid:
|
||||
raise ValueError("Audio list is empty or all elements are None.")
|
||||
return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
|
||||
|
||||
return torch.as_tensor(audio).float().cpu().reshape(-1)
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
engine = Omni(
|
||||
model=args.model,
|
||||
deploy_config=args.deploy_config,
|
||||
)
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import (
|
||||
build_dots_tts_prompt,
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
|
||||
|
||||
prompt = build_dots_tts_prompt(
|
||||
tokenizer=tokenizer,
|
||||
text=args.text,
|
||||
)
|
||||
|
||||
print(f"Model : {args.model}")
|
||||
print(f"Text : {args.text}")
|
||||
print(f"Output dir : {output_dir}")
|
||||
print(f"Prompt len : {len(prompt['prompt_token_ids'])} tokens")
|
||||
|
||||
t_start = time.perf_counter()
|
||||
outputs = engine.generate([prompt])
|
||||
elapsed = time.perf_counter() - t_start
|
||||
|
||||
request_output = outputs[0]
|
||||
mm = request_output.outputs[0].multimodal_output
|
||||
audio = extract_audio(mm)
|
||||
|
||||
duration = audio.numel() / SAMPLE_RATE
|
||||
rtf = elapsed / duration if duration > 0 else float("inf")
|
||||
|
||||
output_path = output_dir / "output.wav"
|
||||
sf.write(str(output_path), audio.numpy(), SAMPLE_RATE, format="WAV")
|
||||
|
||||
print(f"Saved : {output_path}")
|
||||
print(f"Duration : {duration:.2f}s")
|
||||
print(f"Inference : {elapsed:.2f}s")
|
||||
print(f"RTF : {rtf:.3f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
main()
|
||||
@@ -38,17 +38,6 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_MODEL = "fishaudio/s2-pro"
|
||||
DEFAULT_STAGE_CONFIG = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
"..",
|
||||
"..",
|
||||
"..",
|
||||
"..",
|
||||
"vllm_omni",
|
||||
"model_executor",
|
||||
"stage_configs",
|
||||
"fish_speech_s2_pro.yaml",
|
||||
)
|
||||
|
||||
|
||||
def build_prompt(
|
||||
@@ -154,8 +143,6 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
|
||||
)
|
||||
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
|
||||
return omni_kwargs
|
||||
|
||||
|
||||
@@ -278,12 +265,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",
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline inference example for Gepard-1.0 TTS via vLLM-Omni.
|
||||
|
||||
Single-stage native-AR pipeline: a Qwen3.5 backbone samples one 32-code FSQ
|
||||
frame per step; the NeMo NanoCodec decodes the frames to a 22.05 kHz mono
|
||||
waveform. Zero-shot: the default learned voice, no reference audio.
|
||||
|
||||
Generation length and seed are stage settings rather than flags here. One
|
||||
output token is one audio frame, so ``max_tokens`` in the deploy YAML is the
|
||||
frame budget; pass a copy of the YAML via ``--deploy-config`` to change it.
|
||||
|
||||
Usage:
|
||||
python end2end.py --text "Hello, this is Gepard speaking."
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
import vllm_omni
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
|
||||
from vllm_omni import Omni # noqa: E402
|
||||
|
||||
MODEL = "nineninesix/gepard-1.0"
|
||||
SAMPLE_RATE = 22050
|
||||
# Absolute: the loader hands the path straight to OmegaConf, which resolves a
|
||||
# bare name against the current working directory, not the package.
|
||||
DEFAULT_DEPLOY_CONFIG = str(Path(vllm_omni.__file__).resolve().parent / "deploy" / "gepard.yaml")
|
||||
|
||||
|
||||
def build_request(text: str) -> dict:
|
||||
"""Build an Omni request payload for Gepard (zero-shot, no ref audio).
|
||||
|
||||
The model only consumes the [speaker slots, SOT, text, EOT, SOS] layout, so
|
||||
a bare text prompt would have no speaker slots and no SOS.
|
||||
"""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
|
||||
from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids
|
||||
|
||||
# From the checkpoint's own sidecar, which is what the worker builds its
|
||||
# config from too. A layout that disagrees with the trained one does not
|
||||
# fail, it just degrades the speech.
|
||||
cfg = GepardConfig.from_checkpoint(MODEL)
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
|
||||
prompt_token_ids = build_gepard_prompt_ids(
|
||||
tokenizer(text, add_special_tokens=False)["input_ids"],
|
||||
config=cfg,
|
||||
)
|
||||
return {
|
||||
"prompt_token_ids": prompt_token_ids,
|
||||
"additional_information": {"text": [text]},
|
||||
}
|
||||
|
||||
|
||||
def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = SAMPLE_RATE) -> None:
|
||||
sf.write(path, waveform.float().numpy(), sample_rate)
|
||||
seconds = waveform.numel() / sample_rate
|
||||
print(
|
||||
f" Saved {path} ({tuple(waveform.shape)}, {sample_rate} Hz, "
|
||||
f"{seconds:.2f}s, peak {float(waveform.abs().max()):.3f})"
|
||||
)
|
||||
|
||||
|
||||
def main(args) -> None:
|
||||
omni = Omni(
|
||||
model=MODEL,
|
||||
deploy_config=args.deploy_config,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
init_timeout=args.init_timeout,
|
||||
)
|
||||
|
||||
# No explicit SamplingParams: one supplied here would replace the stage
|
||||
# defaults rather than merge over them, dropping the pipeline's stop token.
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"Synthesizing: {args.text!r}")
|
||||
inputs = build_request(text=args.text)
|
||||
|
||||
for i, stage_outputs in enumerate(omni.generate(inputs)):
|
||||
for j, out in enumerate(stage_outputs.outputs):
|
||||
mm = out.multimodal_output
|
||||
if mm is None:
|
||||
print(f" [req {i}] No audio output.")
|
||||
continue
|
||||
# The consolidation path renames "model_outputs" to "audio".
|
||||
audio = mm.get("audio")
|
||||
if audio is None:
|
||||
audio = mm.get("model_outputs")
|
||||
if isinstance(audio, list):
|
||||
audio = audio[0] if len(audio) else None
|
||||
if audio is None:
|
||||
print(f" [req {i}] No waveform in multimodal_output (keys: {list(mm)}).")
|
||||
continue
|
||||
sr_tensor = mm.get("sr")
|
||||
sr = int(sr_tensor.item()) if hasattr(sr_tensor, "item") else SAMPLE_RATE
|
||||
save_audio(audio.cpu(), str(output_dir / f"output_{i}_{j}.wav"), sr)
|
||||
|
||||
print("Done.")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(description="Gepard-1.0 offline TTS inference (zero-shot)")
|
||||
parser.add_argument("--text", default="Hello, this is Gepard speaking.", help="Text to synthesize.")
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default=os.path.join(
|
||||
os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")),
|
||||
"gepard_output",
|
||||
),
|
||||
help="Directory for WAV outputs (default: ~/.cache/gepard_output).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
default=DEFAULT_DEPLOY_CONFIG,
|
||||
help=(
|
||||
"Path to the deploy YAML (default: the packaged vllm_omni/deploy/gepard.yaml). This must "
|
||||
"stay set: without the YAML's `pipeline: gepard` pin the architectures fallback routes "
|
||||
"Qwen3_5ForCausalLM to the diffusion registry. Copy it to change max_tokens or seed."
|
||||
),
|
||||
)
|
||||
# The upstream defaults are too tight for a cold start.
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=900)
|
||||
parser.add_argument("--init-timeout", type=int, default=1800)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(parse_args())
|
||||
@@ -145,7 +145,7 @@ def main(args):
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
stage_configs_path=deploy_config_path,
|
||||
deploy_config=deploy_config_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
@@ -490,10 +490,10 @@ def parse_args():
|
||||
help="Path to a .txt file with one prompt per line (preferred).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a stage configs file.",
|
||||
help="Path to a deploy config YAML.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-path",
|
||||
|
||||
@@ -32,11 +32,10 @@ from end2end import _estimate_prompt_len
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config
|
||||
|
||||
|
||||
def _default_stage_config() -> str:
|
||||
def _default_deploy_config() -> str:
|
||||
repo_root = Path(__file__).resolve().parents[4]
|
||||
return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")
|
||||
|
||||
@@ -109,7 +108,7 @@ def main(args: Any) -> None:
|
||||
"output_dir",
|
||||
):
|
||||
omni_kwargs.pop(key, None)
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
|
||||
omni_kwargs["deploy_config"] = args.deploy_config or _default_deploy_config()
|
||||
omni_kwargs["log_stats"] = args.log_stats
|
||||
omni = Omni(**omni_kwargs)
|
||||
|
||||
@@ -156,7 +155,7 @@ def parse_args() -> Any:
|
||||
default=None,
|
||||
help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
|
||||
)
|
||||
parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML")
|
||||
parser.add_argument("--deploy-config", default=None, help="Qwen3-TTS deploy YAML")
|
||||
parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align")
|
||||
parser.add_argument("--language", default="English", help="Qwen3-TTS language field")
|
||||
parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name")
|
||||
@@ -164,7 +163,6 @@ def parse_args() -> Any:
|
||||
parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens")
|
||||
parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar")
|
||||
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging")
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -570,6 +570,8 @@ def main():
|
||||
num_frames=args.num_frames,
|
||||
extra_args=extra_args,
|
||||
)
|
||||
if args.frame_rate is not None:
|
||||
sampling_kwargs["frame_rate"] = args.frame_rate
|
||||
if args.guidance_scale_high is not None:
|
||||
sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
|
||||
|
||||
|
||||
@@ -0,0 +1,104 @@
|
||||
# π0 (Pi-Zero) VLA — OpenPI realtime serving
|
||||
|
||||
[π0](https://www.physicalintelligence.company/blog/pi0) is a Vision-Language-Action
|
||||
model from Physical Intelligence: multi-camera images + a language instruction +
|
||||
robot proprioceptive state → a continuous action chunk via flow-matching denoising
|
||||
(it does **not** emit text tokens). This example serves π0 over the OpenPI realtime
|
||||
websocket protocol at `/v1/realtime/robot/openpi`.
|
||||
|
||||
## Install extras
|
||||
|
||||
The core `pip install -e .` does not include the OpenPI client used here:
|
||||
|
||||
- `openpi-client` (from the [openpi](https://github.com/Physical-Intelligence/openpi)
|
||||
repo: `pip install -e packages/openpi-client`), `websockets`, `msgpack`, `msgpack-numpy`
|
||||
|
||||
## Weights
|
||||
|
||||
π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the server
|
||||
download it (`MODEL=lerobot/pi0_base`) or point at a local copy
|
||||
(`MODEL=/path/to/pi0_base`).
|
||||
|
||||
## Run the server
|
||||
|
||||
```bash
|
||||
vllm serve lerobot/pi0_base --omni --port 8000 \
|
||||
--served-model-name pi0 \
|
||||
--deploy-config vllm_omni/deploy/pi0.yaml \
|
||||
--enforce-eager --disable-log-stats
|
||||
```
|
||||
|
||||
The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
|
||||
(`Pi0Pipeline`), float32, `max_num_seqs: 1`, and the `policy_server_config` handshake
|
||||
metadata (3 cameras, `joint_position`, action horizon 50, action dim 32).
|
||||
|
||||
### Commonly adjusted `model_config` keys
|
||||
|
||||
These live under `stages[0].model_config` in `vllm_omni/deploy/pi0.yaml`:
|
||||
|
||||
| Key | Default | Meaning |
|
||||
|---|---|---|
|
||||
| `chunk_size` | `50` | Action-chunk length (timesteps) the model predicts per inference. |
|
||||
| `num_inference_steps` | `10` | Flow-matching Euler denoising steps. |
|
||||
| `max_action_dim` | `32` | Action dimensionality (state/action are padded to this). |
|
||||
| `max_state_dim` | `32` | Proprioceptive-state dimensionality (zero-padded to this). |
|
||||
| `image_resolution` | `[224, 224]` | Per-camera input size (square; SigLIP). |
|
||||
| `tokenizer_max_length` | `48` | Max PaliGemma prompt tokens. |
|
||||
| `max_cameras` | `3` | Camera slots the model attends to (real + `-1`-padded). |
|
||||
| `image_feature_keys` | 3 `observation.images.*` keys | Camera order the model attends to. |
|
||||
| `image_key_map` | `{}` | Map raw obs camera keys → `image_feature_keys` (empty = verbatim). |
|
||||
|
||||
The `policy_server_config` block below them is the OpenPI handshake metadata
|
||||
advertised to the client; keep its `action_horizon` / `action_dim` /
|
||||
`image_resolution` in sync with the `model_config` values above. The e2e test
|
||||
`tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online` connects to a
|
||||
live server and asserts the advertised metadata matches these `pi0.yaml` values.
|
||||
|
||||
## Run the client
|
||||
|
||||
```bash
|
||||
python examples/online_serving/pi0/openpi_client.py --host 127.0.0.1 --port 8000 \
|
||||
--prompt "pick up the red block and place it in the bin"
|
||||
```
|
||||
|
||||
It connects, prints the server metadata, sends robot observations, and prints the
|
||||
returned `[action_horizon, action_dim] = [50, 32]` action chunks. Replace the blank
|
||||
cameras / zero state in `_make_dummy_obs` with real frames (HWC uint8) and
|
||||
proprioceptive state to drive a robot.
|
||||
|
||||
### Observation format
|
||||
|
||||
The client sends a flat dict per inference:
|
||||
|
||||
```python
|
||||
{
|
||||
"observation.images.base_0_rgb": np.uint8[H, W, 3],
|
||||
"observation.images.left_wrist_0_rgb": np.uint8[H, W, 3],
|
||||
"observation.images.right_wrist_0_rgb":np.uint8[H, W, 3],
|
||||
"state": np.float32[state_dim], # zero-padded to max_state_dim=32 server-side
|
||||
"prompt": "pick up the red block",
|
||||
"session_id": "<uuid>", # accepted but ignored (π0 is stateless)
|
||||
}
|
||||
```
|
||||
|
||||
Camera keys must match the server's `image_feature_keys` (the checkpoint's
|
||||
`input_features` order). If your robot uses different camera names, set
|
||||
`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
|
||||
|
||||
## Correctness
|
||||
|
||||
π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy` reference
|
||||
(`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
|
||||
`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`) and is
|
||||
version-stable across transformers releases (the version-stability checks in
|
||||
`tests/diffusion/models/pi0/test_pi0_units.py`). An OpenPI websocket online-serving e2e
|
||||
lives in `tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online`.
|
||||
|
||||
## Limitations
|
||||
|
||||
- **Normalization stats**: `lerobot/pi0_base` uses identity normalization, which is
|
||||
fully supported (state passes through, actions are returned in the model's space).
|
||||
Per-dataset `norm_stats` declared in a checkpoint's `config.json` are honored, but
|
||||
stats stored only in LeRobot's `policy_preprocessor.json` companion safetensors are
|
||||
not yet loaded — a fine-tuned checkpoint that relies on those would need that bridge
|
||||
added before its actions are in real-world units.
|
||||
@@ -0,0 +1,119 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Minimal OpenPI client for the π0 (Pi-Zero) VLA policy server.
|
||||
|
||||
Connects to a running π0 server (see README.md), sends a single robot
|
||||
observation (3 cameras + state + language instruction), and prints the returned
|
||||
continuous action chunk ``[action_horizon, action_dim]``.
|
||||
|
||||
The server speaks the OpenPI websocket protocol:
|
||||
connect -> server sends msgpack(policy_server_config metadata)
|
||||
infer -> client sends msgpack(obs), server replies msgpack(actions ndarray)
|
||||
|
||||
Usage:
|
||||
python openpi_client.py --host 127.0.0.1 --port 8000 \
|
||||
--prompt "pick up the red block"
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
# Avoid shadowing the installed ``openpi_client`` package with this dir.
|
||||
example_dir = str(Path(__file__).resolve().parent)
|
||||
removed = sys.path and sys.path[0] == example_dir
|
||||
if removed:
|
||||
sys.path.pop(0)
|
||||
try:
|
||||
from openpi_client.websocket_client_policy import WebsocketClientPolicy
|
||||
finally:
|
||||
if removed:
|
||||
sys.path.insert(0, example_dir)
|
||||
except ImportError as exc: # pragma: no cover - runtime dependency guard
|
||||
raise ImportError("π0 OpenPI example requires `openpi-client`.") from exc
|
||||
|
||||
logger = logging.getLogger("pi0_openpi_client")
|
||||
|
||||
DEFAULT_HOST = "127.0.0.1"
|
||||
DEFAULT_PORT = 8000
|
||||
DEFAULT_PATH = "/v1/realtime/robot/openpi"
|
||||
DEFAULT_PROMPT = "pick up the red block and place it in the bin"
|
||||
|
||||
# π0 / pi0_base camera identities (must match the server's image_feature_keys).
|
||||
CAMERA_KEYS = (
|
||||
"observation.images.base_0_rgb",
|
||||
"observation.images.left_wrist_0_rgb",
|
||||
"observation.images.right_wrist_0_rgb",
|
||||
)
|
||||
STATE_DIM = 32
|
||||
|
||||
|
||||
def _make_dummy_obs(prompt: str, session_id: str, image_size: int = 224) -> dict[str, Any]:
|
||||
"""Build a single π0 observation with blank cameras + zero state.
|
||||
|
||||
Replace the zeros with real camera frames (HWC uint8) and proprioceptive
|
||||
state to drive an actual robot.
|
||||
"""
|
||||
obs: dict[str, Any] = {cam: np.zeros((image_size, image_size, 3), dtype=np.uint8) for cam in CAMERA_KEYS}
|
||||
obs["state"] = np.zeros(STATE_DIM, dtype=np.float32)
|
||||
obs["prompt"] = prompt
|
||||
obs["session_id"] = session_id
|
||||
return obs
|
||||
|
||||
|
||||
def _ws_uri(host: str, port: int, path: str) -> str:
|
||||
"""Build the full websocket URI. WebsocketClientPolicy takes the whole URI as
|
||||
``host`` (it has no separate path arg)."""
|
||||
if host.startswith(("ws://", "wss://")):
|
||||
return host
|
||||
return f"ws://{host}:{port}{path}"
|
||||
|
||||
|
||||
def _as_action_array(response: Any) -> np.ndarray:
|
||||
if isinstance(response, dict) and response.get("type") == "error":
|
||||
raise RuntimeError(f"Inference failed: {response.get('message', response)}")
|
||||
return np.asarray(response, dtype=np.float32)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description="π0 OpenPI client")
|
||||
parser.add_argument("--host", default=DEFAULT_HOST)
|
||||
parser.add_argument("--port", type=int, default=DEFAULT_PORT)
|
||||
parser.add_argument("--path", default=DEFAULT_PATH)
|
||||
parser.add_argument("--prompt", default=DEFAULT_PROMPT)
|
||||
parser.add_argument("--num-steps", type=int, default=2, help="Inference observations to send.")
|
||||
args = parser.parse_args()
|
||||
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
||||
|
||||
policy = WebsocketClientPolicy(host=_ws_uri(args.host, args.port, args.path))
|
||||
metadata = dict(policy.get_server_metadata())
|
||||
logger.info("Server metadata: %s", metadata)
|
||||
|
||||
session_id = str(uuid.uuid4())
|
||||
for step in range(args.num_steps):
|
||||
obs = _make_dummy_obs(args.prompt, session_id)
|
||||
actions = _as_action_array(policy.infer(obs))
|
||||
logger.info(
|
||||
"[step %d] actions shape=%s mean=%.4f std=%.4f",
|
||||
step,
|
||||
actions.shape,
|
||||
float(actions.mean()),
|
||||
float(actions.std()),
|
||||
)
|
||||
if not np.isfinite(actions).all():
|
||||
raise RuntimeError("Server returned non-finite actions.")
|
||||
logger.info("Done — π0 OpenPI round-trip OK.")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -73,12 +73,19 @@ plugins:
|
||||
- awesome-nav
|
||||
- redirects:
|
||||
redirect_maps:
|
||||
community/meetups.md: community/meetings.md
|
||||
community/volunteers.md: community/contact_us.md
|
||||
user_guide/diffusion/request_batching.md: user_guide/diffusion/execution_modes.md
|
||||
user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
|
||||
user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
|
||||
user_guide/quantization/index.md: user_guide/quantization/overview.md
|
||||
configuration/pd_disaggregation.md: features/pd_disaggregation.md
|
||||
design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
|
||||
design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
|
||||
design/feature/cpu_offloading.md: design/feature/offloader/README.md
|
||||
design/feature/model_level_offload.md: design/feature/offloader/cpu_offload.md
|
||||
design/feature/layerwise_offload.md: design/feature/offloader/layerwise_offload.md
|
||||
design/feature/distributed_layerwise_offload.md: design/feature/offloader/distributed_layerwise_offload.md
|
||||
- glightbox
|
||||
- git-revision-date-localized:
|
||||
# exclude files
|
||||
|
||||
Loaded 100 of 545 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user