Compare commits

..
Author SHA1 Message Date
Hongsheng Liu bed9ad54cb Revert "[BugFix] ignore mm data from stages to async omni (#954)"
This reverts commit 6fb4058566.
2026-01-28 17:40:42 +08:00
522 changed files with 116037 additions and 0 deletions

No files matched your search

+238
View File
@@ -0,0 +1,238 @@
#!/bin/bash
# vllm-omni customized version
# Based on: https://github.com/vllm-project/ci-infra/blob/main/buildkite/bootstrap-amd.sh
# Last synced: 2025-12-15
# Modifications: Use local template file instead of downloading from ci-infra
set -euo pipefail
if [[ -z "${RUN_ALL:-}" ]]; then
RUN_ALL=0
fi
if [[ -z "${NIGHTLY:-}" ]]; then
NIGHTLY=0
fi
if [[ -z "${VLLM_CI_BRANCH:-}" ]]; then
VLLM_CI_BRANCH="main"
fi
if [[ -z "${AMD_MIRROR_HW:-}" ]]; then
AMD_MIRROR_HW="amdproduction"
fi
if [[ -z "${DOCS_ONLY_DISABLE:-}" ]]; then
DOCS_ONLY_DISABLE=0
fi
fail_fast() {
DISABLE_LABEL="ci-no-fail-fast"
# If BUILDKITE_PULL_REQUEST != "false", then we check the PR labels using curl and jq
if [ "$BUILDKITE_PULL_REQUEST" != "false" ]; then
PR_LABELS=$(curl -s "https://api.github.com/repos/vllm-project/vllm-omni/pulls/$BUILDKITE_PULL_REQUEST" | jq -r '.labels[].name')
if [[ $PR_LABELS == *"$DISABLE_LABEL"* ]]; then
echo false
else
echo true
fi
else
echo false # not a PR or BUILDKITE_PULL_REQUEST not set
fi
}
check_run_all_label() {
RUN_ALL_LABEL="ready-run-all-tests"
# If BUILDKITE_PULL_REQUEST != "false", then we check the PR labels using curl and jq
if [ "$BUILDKITE_PULL_REQUEST" != "false" ]; then
PR_LABELS=$(curl -s "https://api.github.com/repos/vllm-project/vllm-omni/pulls/$BUILDKITE_PULL_REQUEST" | jq -r '.labels[].name')
if [[ $PR_LABELS == *"$RUN_ALL_LABEL"* ]]; then
echo true
else
echo false
fi
else
echo false # not a PR or BUILDKITE_PULL_REQUEST not set
fi
}
if [[ -z "${COV_ENABLED:-}" ]]; then
COV_ENABLED=0
fi
upload_pipeline() {
echo "Uploading pipeline..."
# Install minijinja
ls .buildkite || buildkite-agent annotate --style error 'Please merge upstream main branch for buildkite CI'
curl -sSfL https://github.com/mitsuhiko/minijinja/releases/download/2.3.1/minijinja-cli-installer.sh | sh
source /var/lib/buildkite-agent/.cargo/env
if [[ $BUILDKITE_PIPELINE_SLUG == "fastcheck" ]]; then
AMD_MIRROR_HW="amdtentative"
fi
# Use local template file for vllm-omni
cp .buildkite/test-template-amd-omni.j2 .buildkite/test-template.j2
# (WIP) Use pipeline generator instead of jinja template
if [ -e ".buildkite/pipeline_generator/pipeline_generator.py" ]; then
python -m pip install click pydantic
python .buildkite/pipeline_generator/pipeline_generator.py --run_all=$RUN_ALL --list_file_diff="$LIST_FILE_DIFF" --nightly="$NIGHTLY" --mirror_hw="$AMD_MIRROR_HW"
buildkite-agent pipeline upload .buildkite/pipeline.yaml
exit 0
fi
echo "List file diff: $LIST_FILE_DIFF"
echo "Run all: $RUN_ALL"
echo "Nightly: $NIGHTLY"
echo "AMD Mirror HW: $AMD_MIRROR_HW"
FAIL_FAST=$(fail_fast)
cd .buildkite
(
set -x
# Output pipeline.yaml with all blank lines removed
minijinja-cli test-template.j2 test-amd.yaml \
-D branch="$BUILDKITE_BRANCH" \
-D list_file_diff="$LIST_FILE_DIFF" \
-D run_all="$RUN_ALL" \
-D nightly="$NIGHTLY" \
-D mirror_hw="$AMD_MIRROR_HW" \
-D fail_fast="$FAIL_FAST" \
-D vllm_use_precompiled="$VLLM_USE_PRECOMPILED" \
-D vllm_merge_base_commit="$(git merge-base origin/main HEAD)" \
-D cov_enabled="$COV_ENABLED" \
-D vllm_ci_branch="$VLLM_CI_BRANCH" \
| sed '/^[[:space:]]*$/d' \
> pipeline.yaml
)
cat pipeline.yaml
buildkite-agent artifact upload pipeline.yaml
buildkite-agent pipeline upload pipeline.yaml
exit 0
}
get_diff() {
$(git add .)
echo $(git diff --name-only --diff-filter=ACMDR $(git merge-base origin/main HEAD))
}
get_diff_main() {
$(git add .)
echo $(git diff --name-only --diff-filter=ACMDR HEAD~1)
}
file_diff=$(get_diff)
if [[ $BUILDKITE_BRANCH == "main" ]]; then
file_diff=$(get_diff_main)
fi
# ----------------------------------------------------------------------
# Early exit start: skip pipeline if conditions are met
# ----------------------------------------------------------------------
# skip pipeline if all changed files are under docs/
if [[ "${DOCS_ONLY_DISABLE}" != "1" ]]; then
if [[ -n "${file_diff:-}" ]]; then
docs_only=1
# Robust iteration over newline-separated file_diff
while IFS= read -r f; do
[[ -z "$f" ]] && continue
# **Policy:** only skip if *every* path starts with docs/
if [[ "$f" != docs/* ]]; then
docs_only=0
break
fi
done < <(printf '%s\n' "$file_diff" | tr ' ' '\n' | tr -d '\r')
if [[ "$docs_only" -eq 1 ]]; then
buildkite-agent annotate ":memo: CI skipped — docs/** only changes detected
\`\`\`
${file_diff}
\`\`\`" --style "info" || true
echo "[docs-only] All changes are under docs/. Exiting before pipeline upload."
exit 0
fi
fi
fi
# ----------------------------------------------------------------------
# Early exit end
# ----------------------------------------------------------------------
patterns=(
"docker/Dockerfile"
"CMakeLists.txt"
"requirements/common.txt"
"requirements/cuda.txt"
"requirements/build.txt"
"requirements/test.txt"
"setup.py"
"csrc/"
"cmake/"
)
ignore_patterns=(
"docker/Dockerfile."
"csrc/cpu"
"csrc/rocm"
"cmake/hipify.py"
"cmake/cpu_extension.cmake"
)
for file in $file_diff; do
# First check if file matches any pattern
matches_pattern=0
for pattern in "${patterns[@]}"; do
if [[ $file == $pattern* ]] || [[ $file == $pattern ]]; then
matches_pattern=1
break
fi
done
# If file matches pattern, check it's not in ignore patterns
if [[ $matches_pattern -eq 1 ]]; then
matches_ignore=0
for ignore in "${ignore_patterns[@]}"; do
if [[ $file == $ignore* ]] || [[ $file == $ignore ]]; then
matches_ignore=1
break
fi
done
if [[ $matches_ignore -eq 0 ]]; then
RUN_ALL=1
echo "Found changes: $file. Run all tests"
break
fi
fi
done
# Check for ready-run-all-tests label
LABEL_RUN_ALL=$(check_run_all_label)
if [[ $LABEL_RUN_ALL == true ]]; then
RUN_ALL=1
NIGHTLY=1
echo "Found 'ready-run-all-tests' label. Running all tests including optional tests."
fi
# Decide whether to use precompiled wheels
# Relies on existing patterns array as a basis.
if [[ -n "${VLLM_USE_PRECOMPILED:-}" ]]; then
echo "VLLM_USE_PRECOMPILED is already set to: $VLLM_USE_PRECOMPILED"
elif [[ $RUN_ALL -eq 1 ]]; then
export VLLM_USE_PRECOMPILED=0
echo "Detected critical changes, building wheels from source"
else
export VLLM_USE_PRECOMPILED=1
echo "No critical changes, using precompiled wheels"
fi
LIST_FILE_DIFF=$(get_diff | tr ' ' '|')
if [[ $BUILDKITE_BRANCH == "main" ]]; then
LIST_FILE_DIFF=$(get_diff_main | tr ' ' '|')
fi
upload_pipeline
+256
View File
@@ -0,0 +1,256 @@
steps:
- label: ":docker: Build image"
key: image-build
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
agents:
queue: "cpu_queue_premerge"
# - label: "Test on NPU"
# depends_on: ~
# key: npu-test
# commands:
# - ".buildkite/scripts/hardware_ci/run_npu_test.sh"
# agents:
# queue: "ascend"
- label: "Simple Unit Test"
depends_on: ~
commands:
- ".buildkite/scripts/simple_test.sh"
agents:
queue: "cpu_queue_premerge"
- label: "Diffusion Model Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Images API LoRA E2E"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Model CPU offloading Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Audio Generation Model Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Cache Backend Test"
timeout_in_minutes: 15
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Sequence Parallelism Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Tensor Parallelism Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion GPU Worker Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test"
timeout_in_minutes: 15
depends_on: image-build
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test with H100"
timeout_in_minutes: 30
depends_on: image-build
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
- pytest -s -v tests/e2e/online_serving/test_async_omni.py
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Diffusion Image Edit Test with H100 (1 GPU)"
timeout_in_minutes: 20
depends_on: image-build
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_i2i_multi_image_input.py
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
+73
View File
@@ -0,0 +1,73 @@
#!/bin/bash
# Helper function to safely login to ECR Public with per-job config isolation
# Uses DOCKER_CONFIG environment variable to prevent race conditions
#
# This script prevents the "device or resource busy" error by giving each
# Buildkite job its own isolated Docker config directory.
#
# Usage:
# source docker_login_ecr_public.sh && safe_docker_login_ecr_public
set -euo pipefail
# Configuration
ECR_REGISTRY="public.ecr.aws"
setup_isolated_docker_config() {
# Use BUILDKITE_JOB_ID for job-specific isolation
# Fallback to PID if running outside Buildkite
local job_id="${BUILDKITE_JOB_ID:-$$}"
# Set Docker config to job-specific directory
export DOCKER_CONFIG="/tmp/docker-config-${job_id}"
# Create directory if it doesn't exist
mkdir -p "$DOCKER_CONFIG"
echo "[docker-config] Using isolated Docker config: $DOCKER_CONFIG"
}
check_docker_auth() {
# Check if already authenticated to the given registry
# Returns 0 if authenticated, 1 if not
local registry="$1"
# Check if credentials exist in the isolated config
if [[ -f "$DOCKER_CONFIG/config.json" ]]; then
# Check if registry is present in config
if grep -q "$registry" "$DOCKER_CONFIG/config.json" 2>/dev/null; then
return 0
fi
fi
return 1
}
safe_docker_login_ecr_public() {
# Setup isolated config first
setup_isolated_docker_config
local registry="$ECR_REGISTRY"
# Check if already authenticated (within this job)
if check_docker_auth "$registry"; then
echo "[docker-login] Already authenticated to $registry in this job"
return 0
fi
# Perform login to isolated config directory
echo "[docker-login] Logging in to $ECR_REGISTRY (isolated config)..."
if aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$ECR_REGISTRY"; then
echo "[docker-login] Login successful (config: $DOCKER_CONFIG)"
return 0
else
local exit_code=$?
echo "[docker-login] ERROR: Login failed with exit code $exit_code" >&2
return $exit_code
fi
}
# Execute if run as script (not sourced)
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
safe_docker_login_ecr_public
fi
+160
View File
@@ -0,0 +1,160 @@
#!/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
# Export Python path
export PYTHONPATH=".."
# Print ROCm version
echo "--- Confirming Clean Initial State"
while true; do
sleep 3
if grep -q clean /opt/amdgpu/etc/gpu_state; then
echo "GPUs state is \"clean\""
break
fi
done
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
}
# Call the cleanup docker function
cleanup_docker
echo "--- Resetting GPUs"
echo "reset" > /opt/amdgpu/etc/gpu_state
while true; do
sleep 3
if grep -q clean /opt/amdgpu/etc/gpu_state; then
echo "GPUs state is \"clean\""
break
fi
done
echo "--- Pulling container"
image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
# 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]}")/../docker_login_ecr_public.sh"
safe_docker_login_ecr_public
# Pull the container from ECR Public Gallery
docker pull "${image_name}"
remove_docker_container() {
docker rm -f "${container_name}" || 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"
commands=$@
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.
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" \
--rm \
-e HIP_VISIBLE_DEVICES="${GPU}" \
-e HF_TOKEN \
-e AWS_ACCESS_KEY_ID \
-e AWS_SECRET_ACCESS_KEY \
-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}"
fi
done
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" \
--rm \
-e HF_TOKEN \
-e AWS_ACCESS_KEY_ID \
-e AWS_SECRET_ACCESS_KEY \
-v "${HF_CACHE}:${HF_MOUNT}" \
-e "HF_HOME=${HF_MOUNT}" \
-e "PYTHONPATH=${MYPYTHONPATH}" \
--name "${container_name}" \
"${image_name}" \
/bin/bash -c "${commands}"
fi
@@ -0,0 +1,145 @@
#!/bin/bash
# This script build the Ascend NPU docker image and run the offline inference inside the container.
# It serves a sanity check for compilation and basic model usage.
set -ex
# Base ubuntu image with basic ascend development libraries and python installed
VLLM_OMNI_REPO="https://github.com/vllm-project/vllm-omni.git"
BASE_IMAGE_NAME="quay.nju.edu.cn/ascend/vllm-ascend:v0.11.0rc2"
image_name="npu/vllm-omni-ci:${BUILDKITE_COMMIT}_${EPOCHSECONDS}"
# image_name="npu/vllm-ci:${BUILDKITE_COMMIT}_${EPOCHSECONDS}"
container_name="npu_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
# BUILDKITE_AGENT_NAME format is {hostname}-{agent_idx}-{npu_card_num}cards
agent_idx=$(echo "${BUILDKITE_AGENT_NAME}" | awk -F'-' '{print $(NF-1)}')
echo "agent_idx: ${agent_idx}"
builder_name="cachebuilder${agent_idx}"
builder_cache_dir="/mnt/docker-cache${agent_idx}"
mkdir -p ${builder_cache_dir}
# Try building the docker image
cat <<EOF | DOCKER_BUILDKIT=1 docker build \
--add-host pypi-cache:${PYPI_CACHE_HOST} \
--builder ${builder_name} --cache-from type=local,src=${builder_cache_dir} \
--cache-to type=local,dest=${builder_cache_dir},mode=max \
--build-arg BUILDKITE_PULL_REQUEST="${BUILDKITE_PULL_REQUEST}" \
--build-arg BUILDKITE_PULL_REQUEST_REPO="${BUILDKITE_PULL_REQUEST_REPO}" \
--progress=plain --load -t ${image_name} -f - .
FROM ${BASE_IMAGE_NAME}
# Define environments
ENV DEBIAN_FRONTEND=noninteractive
RUN pip config set global.index-url http://pypi-cache:${PYPI_CACHE_PORT}/pypi/simple && \
pip config set global.trusted-host pypi-cache && \
apt-get update -y && \
apt-get install -y python3-pip git vim wget net-tools gcc g++ cmake libnuma-dev && \
rm -rf /var/cache/apt/* && \
rm -rf /var/lib/apt/lists/*
# Install for pytest to make the docker build cache layer always valid
RUN --mount=type=cache,target=/root/.cache/pip \
pip install pytest>=6.0 pytest-cov modelscope
COPY . .
# Install vllm-omni
WORKDIR /workspace
ARG VLLM_OMNI_REPO=https://github.com/vllm-project/vllm-omni.git
ARG VLLM_OMNI_TAG=main
ARG BUILDKITE_PULL_REQUEST
ARG BUILDKITE_PULL_REQUEST_REPO
RUN git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf "https://github.com/" && \
if [ "\$BUILDKITE_PULL_REQUEST" != "false" ] && [ -n "\$BUILDKITE_PULL_REQUEST" ]; then \
echo "Cloning and checking out PR #\$BUILDKITE_PULL_REQUEST..." && \
git clone \$VLLM_OMNI_REPO /workspace/vllm-omni && \
cd /workspace/vllm-omni && \
git fetch origin pull/\$BUILDKITE_PULL_REQUEST/head:pr-\$BUILDKITE_PULL_REQUEST && \
git checkout pr-\$BUILDKITE_PULL_REQUEST; \
else \
echo "Not a PR build, using main branch" && \
git clone --depth 1 \$VLLM_OMNI_REPO /workspace/vllm-omni; \
fi
RUN --mount=type=cache,target=/root/.cache/pip \
export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
source /usr/local/Ascend/nnal/atb/set_env.sh && \
export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
python3 -m pip install -v -e /workspace/vllm-omni/
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
WORKDIR /workspace/vllm-omni
CMD ["/bin/bash"]
EOF
# Setup cleanup
remove_docker_container() {
docker rm -f "${container_name}" || true;
docker image rm -f "${image_name}" || true;
docker system prune -f || true;
}
trap remove_docker_container EXIT
# Generate corresponding --device args based on BUILDKITE_AGENT_NAME
# Ascend NPU BUILDKITE_AGENT_NAME format is {hostname}-{agent_idx}-{npu_card_num}cards, and agent_idx starts from 1.
# e.g. atlas-a2-001-1-2cards means this is the 1-th agent on atlas-a2-001 host, and it has 2 NPU cards.
# returns --device /dev/davinci0 --device /dev/davinci1
parse_and_gen_devices() {
local input="$1"
local index cards_num
if [[ "$input" =~ ([0-9]+)-([0-9]+)cards$ ]]; then
index="${BASH_REMATCH[1]}"
cards_num="${BASH_REMATCH[2]}"
else
echo "parse error" >&2
return 1
fi
local devices=""
local i=0
while (( i < cards_num )); do
local dev_idx=$(((index - 1)*cards_num + i ))
devices="$devices --device /dev/davinci${dev_idx}"
((i++))
done
# trim leading space
devices="${devices#"${devices%%[![:space:]]*}"}"
# Output devices: assigned to the caller variable
printf '%s' "$devices"
}
devices=$(parse_and_gen_devices "${BUILDKITE_AGENT_NAME}") || exit 1
# Run the image and execute the Out-Of-Tree (OOT) platform interface test case on Ascend NPU hardware.
# This test checks whether the OOT platform interface is functioning properly in conjunction with
# the hardware plugin vllm-ascend.
hf_model_cache_dir=/mnt/hf_cache${agent_idx}
ms_model_cache_dir=/mnt/modelscope${agent_idx}
mkdir -p ${hf_model_cache_dir}
mkdir -p ${ms_model_cache_dir}
docker run \
--init \
${devices} \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
-v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v ${hf_model_cache_dir}:/root/.cache/huggingface \
-v ${ms_model_cache_dir}:/root/.cache/modelscope \
--network host \
--entrypoint="" \
--name "${container_name}" \
"${image_name}" \
bash -c '
set -e
VLLM_USE_MODELSCOPE=True pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
'
+58
View File
@@ -0,0 +1,58 @@
#!/usr/bin/env bash
set -euo pipefail
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)"
UV_BIN="uv"
PYTHON_BOOTSTRAP="${PYTHON:-python3}"
VENV_DIR="${VENV_DIR:-${REPO_ROOT}/.venv-simple-test}"
TARGET_PY_VERSION="3.12"
if ! command -v "${PYTHON_BOOTSTRAP}" >/dev/null 2>&1; then
PYTHON_BOOTSTRAP="python"
fi
cd "${REPO_ROOT}"
# Ensure pip-installed scripts (including uv) are on PATH
PYTHON_SCRIPTS="$(${PYTHON_BOOTSTRAP} -c 'import sysconfig; print(sysconfig.get_path("scripts"))')"
USER_BASE="$(${PYTHON_BOOTSTRAP} -m site --user-base 2>/dev/null || true)"
PATH="${PYTHON_SCRIPTS}:${PATH}"
if [[ -n "${USER_BASE}" ]]; then
PATH="${USER_BASE}/bin:${PATH}"
fi
if ! command -v "${UV_BIN}" >/dev/null 2>&1; then
"${PYTHON_BOOTSTRAP}" -m pip install --upgrade pip
"${PYTHON_BOOTSTRAP}" -m pip install uv
hash -r
fi
TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
if [[ -z "${TARGET_PY_PATH}" ]]; then
"${UV_BIN}" python install "${TARGET_PY_VERSION}"
TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
fi
if [[ -z "${TARGET_PY_PATH}" ]]; then
echo "Failed to locate python ${TARGET_PY_VERSION}" >&2
exit 1
fi
if [[ ! -d "${VENV_DIR}" ]]; then
"${UV_BIN}" venv --python "${TARGET_PY_PATH}" "${VENV_DIR}"
fi
VENV_PYTHON="${VENV_DIR}/bin/python"
[[ -x "${VENV_PYTHON}" ]] || { echo "Python not found in ${VENV_DIR}"; exit 1; }
"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.14.0
"${UV_BIN}" pip install --python "${VENV_PYTHON}" -e ".[dev]"
"${VENV_PYTHON}" -m pytest -v -s tests/entrypoints/
"${VENV_PYTHON}" -m pytest -v -s tests/diffusion/cache/
"${VENV_PYTHON}" -m pytest -v -s tests/diffusion/lora/
"${VENV_PYTHON}" -m pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
"${VENV_PYTHON}" -m pytest -v -s tests/worker/
"${VENV_PYTHON}" -m pytest -v -s tests/distributed/omni_connectors/test_kv_flow.py
+106
View File
@@ -0,0 +1,106 @@
steps:
- label: "Diffusion Model Test"
timeout_in_minutes: 20
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export MIOPEN_DEBUG_CONV_DIRECT=0
- export MIOPEN_DEBUG_CONV_GEMM=0
- export VLLM_ROCM_USE_AITER=1
- export VLLM_ROCM_USE_AITER_MHA=1
- export VLLM_ROCM_USE_AITER_LINEAR=0
- export VLLM_ROCM_USE_AITER_RMSNORM=0
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
- label: "Diffusion Cache Backend Test"
timeout_in_minutes: 15
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export MIOPEN_DEBUG_CONV_DIRECT=0
- export MIOPEN_DEBUG_CONV_GEMM=0
- export VLLM_ROCM_USE_AITER=1
- export VLLM_ROCM_USE_AITER_MHA=1
- export VLLM_ROCM_USE_AITER_LINEAR=0
- export VLLM_ROCM_USE_AITER_RMSNORM=0
- pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
- label: "Diffusion Parallelism Test"
timeout_in_minutes: 15
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export MIOPEN_DEBUG_CONV_DIRECT=0
- export MIOPEN_DEBUG_CONV_GEMM=0
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
- label: "Diffusion GPU Worker Test"
timeout_in_minutes: 20
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export MIOPEN_DEBUG_CONV_DIRECT=0
- export MIOPEN_DEBUG_CONV_GEMM=0
- pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
- label: "Omni Model Test Qwen2-5-Omni"
timeout_in_minutes: 15
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export MIOPEN_DEBUG_CONV_DIRECT=0
- export MIOPEN_DEBUG_CONV_GEMM=0
- export VLLM_ROCM_USE_AITER=1
- export VLLM_ROCM_USE_AITER_MHA=1
- export VLLM_ROCM_USE_AITER_LINEAR=0
- export VLLM_ROCM_USE_AITER_RMSNORM=0
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
- label: "Omni Model Test Qwen3-Omni"
timeout_in_minutes: 15
agent_pool: mi325_2
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export MIOPEN_DEBUG_CONV_DIRECT=0
- export MIOPEN_DEBUG_CONV_GEMM=0
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py
- label: "Diffusion Image Edit Test"
timeout_in_minutes: 15
agent_pool: mi325_1
depends_on: amd-build
mirror_hardwares: [amdproduction]
grade: Blocking
commands:
- export GPU_ARCHS=gfx942
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export MIOPEN_DEBUG_CONV_DIRECT=0
- export MIOPEN_DEBUG_CONV_GEMM=0
- export VLLM_ROCM_USE_AITER=1
- export VLLM_ROCM_USE_AITER_MHA=1
- export VLLM_ROCM_USE_AITER_LINEAR=0
- export VLLM_ROCM_USE_AITER_RMSNORM=0
- pytest -s -v tests/e2e/online_serving/test_i2i_multi_image_input.py
+53
View File
@@ -0,0 +1,53 @@
{# 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
#}
{% set docker_image_amd = "public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT-rocm-omni" %}
{% set default_working_dir = "/app/vllm-omni" %}
- group: "AMD Tests"
depends_on: ~
steps:
- label: "AMD: :docker: build image"
depends_on: ~
soft_fail: false
commands:
- "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
- "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target final --progress plain ."
- "docker push {{ docker_image_amd }}"
key: "amd-build"
env:
DOCKER_BUILDKIT: "1"
retry:
automatic:
- exit_status: -1 # Agent was lost
limit: 1
- exit_status: -10 # Agent was lost
limit: 1
- exit_status: 1 # Machine occasionally fail
limit: 1
agents:
queue: cpu_queue_premerge
{% for step in steps %}
{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
- label: "{{ step.agent_pool }}: {{ step.label }}"
depends_on: amd-build
agents:
{% if step.agent_pool %}
queue: amd_{{ step.agent_pool }}
{% else %}
queue: amd_mi325_1
{% endif %}
command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh "(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }} ; {{ step.command or (step.commands | join(" && ")) | safe }}"
env:
DOCKER_BUILDKIT: "1"
priority: 100
{% if step.grade and step.grade == "Blocking" %}
soft_fail: false
{% else %}
soft_fail: true
{% endif%}
{% endif %}
{% endfor %}
+44
View File
@@ -0,0 +1,44 @@
# Set default behavior to automatically normalize line endings
* text=auto
# Explicitly declare text files you want to always be normalized and converted
# to native line endings on checkout
*.py text eol=lf
*.yaml text eol=lf
*.yml text eol=lf
*.md text eol=lf
*.txt text eol=lf
*.json text eol=lf
*.toml text eol=lf
*.sh text eol=lf
*.bat text eol=crlf
*.cmd text eol=crlf
# Declare files that will always have CRLF line endings on checkout
*.bat text eol=crlf
*.cmd text eol=crlf
# Denote all files that are truly binary and should not be modified
*.png binary
*.jpg binary
*.jpeg binary
*.gif binary
*.ico binary
*.mov binary
*.mp4 binary
*.mp3 binary
*.flv binary
*.fla binary
*.swf binary
*.gz binary
*.zip binary
*.7z binary
*.ttf binary
*.eot binary
*.woff binary
*.woff2 binary
*.pyc binary
*.pdf binary
*.ai binary
*.psd binary
*.xcf binary
+17
View File
@@ -0,0 +1,17 @@
# CODEOWNERS is used to automatically assign PR reviewers
# Format: path pattern @username or @team
# More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
# Default owners (fallback reviewers for all files)
# later, we can add more reviewers here
* @hsliuustc0106
# Assign reviewers by directory
/.github/ @owner1
/tests/ @owner1
/docs/ @owner1
# Example: reviewers for specific directories
# /examples/offline_inference/ @offline-inference-team
# /vllm_omni/core/ @core-team
# /vllm_omni/model_executor/ @model-team
@@ -0,0 +1,29 @@
name: 📚 Documentation
description: Report an issue related to documentation
title: "[Doc]: "
labels: ["documentation"]
body:
- type: textarea
attributes:
label: 📚 The doc issue
description: >
A clear and concise description of what content in the documentation is an issue.
validations:
required: true
- type: textarea
attributes:
label: Suggest a potential alternative/fix
description: >
Tell us how we could improve the documentation in this regard.
- type: markdown
attributes:
value: >
Thanks for contributing 🎉!
- type: checkboxes
id: askllm
attributes:
label: Before submitting a new issue...
options:
- label: Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://vllm-omni.readthedocs.io), which can answer lots of frequently asked questions.
required: true
@@ -0,0 +1,47 @@
name: 🛠️ Installation
description: Report an issue here when you hit errors during installation.
title: "[Installation]: "
labels: ["installation"]
body:
- type: markdown
attributes:
value: >
#### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
- type: textarea
attributes:
label: Your current environment
description: |
Please run the following and paste the output below.
```sh
wget https://raw.githubusercontent.com/vllm-project/vllm-omni/main/vllm_omni/collect_env.py
# For security purposes, please feel free to check the contents of collect_env.py before running it.
python collect_env.py
```
It is suggested to download and execute the latest script, as vllm-omni might frequently update the diagnosis information needed for accurately and quickly responding to issues.
value: |
```text
The output of `python collect_env.py`
```
validations:
required: true
- type: textarea
attributes:
label: How you are installing vllm-omni
description: |
Paste the full command you are trying to execute.
value: |
```sh
pip install -vvv vllm-omni
```
- type: markdown
attributes:
value: >
Thanks for contributing 🎉!
- type: checkboxes
id: askllm
attributes:
label: Before submitting a new issue...
options:
- label: Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://vllm-omni.readthedocs.io), which can answer lots of frequently asked questions.
required: true
+136
View File
@@ -0,0 +1,136 @@
name: 🐛 Bug report
description: Raise an issue here if you find a bug.
title: "[Bug]: "
labels: ["bug"]
body:
- type: markdown
attributes:
value: >
#### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
- type: markdown
attributes:
value: |
⚠️ **SECURITY WARNING:** Please review any text you paste to ensure it does not contain sensitive information such as:
- API tokens or keys (e.g., Hugging Face tokens, OpenAI API keys)
- Passwords or authentication credentials
- Private URLs or endpoints
- Personal or confidential data
Consider redacting or replacing sensitive values with placeholders like `<YOUR_TOKEN_HERE>` when sharing configuration or code examples.
- type: textarea
attributes:
label: Your current environment
description: |
Please run the following and paste the output below.
```sh
wget https://raw.githubusercontent.com/vllm-project/vllm-omni/main/collect_env.py
# For security purposes, please feel free to check the contents of collect_env.py before running it.
python collect_env.py
```
It is suggested to download and execute the latest script, as vllm-omni might frequently update the diagnosis information needed for accurately and quickly responding to issues.
value: |
<details>
<summary>The output of <code>python collect_env.py</code></summary>
```text
Your output of `python collect_env.py` here
```
</details>
validations:
required: true
- type: textarea
attributes:
label: Your code version
description: |
Please provide a description of which commit id or version the code is.
value: |
<details>
<summary>The commit id or version of vllm</summary>
```text
```
</details>
<details>
<summary>The commit id or version of vllm-omni</summary>
```text
```
</details>
validations:
required: true
- type: textarea
attributes:
label: 🐛 Describe the bug
description: |
Please provide a clear and concise description of what the bug is.
If relevant, add a minimal example so that we can reproduce the error by running the code. It is very important for the snippet to be as succinct (minimal) as possible, so please take time to trim down any irrelevant code to help us debug efficiently. We are going to copy-paste your code and we expect to get the same result as you did: avoid any external data, and include the relevant imports, etc. For example:
```python
from vllm_omni import OmniLLM, create_ar_stage_config, create_dit_stage_config
# Create stage configurations
ar_config = create_ar_stage_config(
stage_id=0,
model_path="Qwen/Qwen3-0.6B",
input_modalities=["text"],
output_modalities=["text"]
)
dit_config = create_dit_stage_config(
stage_id=1,
model_path="stabilityai/stable-diffusion-2-1",
input_modalities=["text"],
output_modalities=["image"]
)
# Initialize OmniLLM
omni = OmniLLM([ar_config, dit_config])
# Generate
outputs = omni.generate(prompt="A scenic watercolor painting of a lighthouse at sunset")
```
If the code is too long (hopefully, it isn't), feel free to put it in a public gist and link it in the issue: https://gist.github.com.
Please also paste or describe the results you observe instead of the expected results. If you observe an error, please paste the error message including the **full** traceback of the exception. It may be relevant to wrap error messages in ```` ```triple quotes blocks``` ````.
Please set the environment variable `export VLLM_OMNI_LOGGING_LEVEL=DEBUG` to turn on more logging to help debugging potential issues.
If you experienced crashes or hangs, it would be helpful to run vllm-omni with `export VLLM_OMNI_TRACE_FUNCTION=1` . All the function calls in vllm-omni will be recorded. Inspect these log files, and tell which function crashes or hangs.
placeholder: |
A clear and concise description of what the bug is.
```python
# Sample code to reproduce the problem
```
```
The error message you got, with the full traceback and the error logs with [dump_input.py:##] if present.
```
validations:
required: true
- type: markdown
attributes:
value: |
⚠️ Please separate bugs of `transformers` implementation or usage from bugs of `vllm-omni`. If you think anything is wrong with the model's output:
- Try the counterpart of `transformers` first. If the error appears, please go to [their issues](https://github.com/huggingface/transformers/issues?q=is%3Aissue+is%3Aopen+sort%3Aupdated-desc).
- If the error only appears in vllm-omni, please provide the detailed script of how you run `transformers` and `vllm-omni`, also highlight the difference and what you expect.
Thanks for reporting 🙏!
- type: checkboxes
id: askllm
attributes:
label: Before submitting a new issue...
options:
- label: Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://vllm-omni.readthedocs.io), which can answer lots of frequently asked questions.
required: true
@@ -0,0 +1,38 @@
name: 🚀 Feature request
description: Submit a proposal/request for a new vllm-omni feature
title: "[Feature]: "
labels: ["feature request"]
body:
- type: markdown
attributes:
value: >
#### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
- type: textarea
attributes:
label: 🚀 The feature, motivation and pitch
description: >
A clear and concise description of the feature proposal. Please outline the motivation for the proposal. Is your feature request related to a specific problem? e.g., *"I'm working on X and would like Y to be possible"*. If this is related to another GitHub issue, please link here too.
validations:
required: true
- type: textarea
attributes:
label: Alternatives
description: >
A description of any alternative solutions or features you've considered, if any.
- type: textarea
attributes:
label: Additional context
description: >
Add any other context or screenshots about the feature request.
- type: markdown
attributes:
value: >
Thanks for contributing 🎉!
- type: checkboxes
id: askllm
attributes:
label: Before submitting a new issue...
options:
- label: Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://vllm-omni.readthedocs.io), which can answer lots of frequently asked questions.
required: true
+45
View File
@@ -0,0 +1,45 @@
name: 🤗 Support request for a new model
description: Submit a proposal/request for a new model support
title: "[New Model]: "
labels: ["new model"]
body:
- type: markdown
attributes:
value: >
#### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
#### We also highly recommend you read the contributing guidelines first to understand how to add a new model.
- type: textarea
attributes:
label: The model to consider.
description: >
A huggingface url, pointing to the model, e.g. https://huggingface.co/openai-community/gpt2 .
validations:
required: true
- type: textarea
attributes:
label: The closest model vllm-omni already supports.
description: >
Here is the list of models already supported by vllm-omni. Which model is the most similar to the model you want to add support for?
- type: textarea
attributes:
label: What's your difficulty of supporting the model you want?
description: >
For example, any new operators or new architecture? Is it a multimodal model? Does it require special preprocessing?
- type: textarea
attributes:
label: Use case and motivation
description: >
What is the primary use case for this model? Why is it important to support in vllm-omni?
- type: markdown
attributes:
value: >
Thanks for contributing 🎉!
- type: checkboxes
id: askllm
attributes:
label: Before submitting a new issue...
options:
- label: Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://vllm-omni.readthedocs.io), which can answer lots of frequently asked questions.
required: true
@@ -0,0 +1,59 @@
name: ⚡ Discussion on the performance of vllm-omni
description: Submit a proposal/discussion about the performance of vllm-omni
title: "[Performance]: "
labels: ["performance"]
body:
- type: markdown
attributes:
value: >
#### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm-omni/issues?q=is%3Aissue+sort%3Acreated-desc+).
- type: textarea
attributes:
label: Proposal to improve performance
description: >
How do you plan to improve vllm-omni's performance?
validations:
required: false
- type: textarea
attributes:
label: Report of performance regression
description: >
Please provide detailed description of performance comparison to confirm the regression. You may want to run the benchmark script at https://github.com/vllm-project/vllm-omni/tree/main/tests/benchmarks .
validations:
required: false
- type: textarea
attributes:
label: Misc discussion on performance
description: >
Anything about the performance.
validations:
required: false
- type: textarea
attributes:
label: Your current environment (if you think it is necessary)
description: |
Please run the following and paste the output below.
```sh
wget https://raw.githubusercontent.com/vllm-project/vllm-omni/main/vllm_omni/collect_env.py
# For security purposes, please feel free to check the contents of collect_env.py before running it.
python collect_env.py
```
It is suggested to download and execute the latest script, as vllm-omni might frequently update the diagnosis information needed for accurately and quickly responding to issues.
value: |
```text
The output of `python collect_env.py`
```
validations:
required: false
- type: markdown
attributes:
value: >
Thanks for contributing 🎉!
- type: checkboxes
id: askllm
attributes:
label: Before submitting a new issue...
options:
- label: Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://vllm-omni.readthedocs.io), which can answer lots of frequently asked questions.
required: true
+54
View File
@@ -0,0 +1,54 @@
name: 💬 Request for comments (RFC).
description: Ask for feedback on major architectural changes or design choices.
title: "[RFC]: "
labels: ["RFC"]
body:
- type: markdown
attributes:
value: >
#### Please take a look at previous [RFCs](https://github.com/vllm-project/vllm-omni/issues?q=label%3ARFC+sort%3Aupdated-desc) for reference.
- type: textarea
attributes:
label: Motivation.
description: >
The motivation of the RFC.
validations:
required: true
- type: textarea
attributes:
label: Proposed Change.
description: >
The proposed change of the RFC.
value: |
Please provide the detailed design document of the RFC using the [template](https://docs.google.com/document/d/12YxSsVeD1jvL-InClkeAEnZyWFDndz_65JmXvsamuV4/edit?tab=t.0#heading=h.4ef4szrsgspp).
validations:
required: true
- type: textarea
attributes:
label: Feedback Period.
description: >
The feedback period of the RFC. Usually at least one week.
validations:
required: false
- type: textarea
attributes:
label: CC List.
description: >
The list of people you want to CC.
validations:
required: false
- type: textarea
attributes:
label: Any Other Things.
description: >
Any other things you would like to mention.
validations:
required: false
- type: checkboxes
id: askllm
attributes:
label: Before submitting a new issue...
options:
- label: Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://vllm-omni.readthedocs.io), which can answer lots of frequently asked questions.
required: true
+5
View File
@@ -0,0 +1,5 @@
blank_issues_enabled: false
contact_links:
- name: Questions
url: https://slack.vllm.ai
about: Ask questions and discuss with other vLLM-Omni community members in our #sig-omni slack channel.
+21
View File
@@ -0,0 +1,21 @@
<!-- markdownlint-disable -->
PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
## Purpose
## Test Plan
## Test Result
---
<details>
<summary> Essential Elements of an Effective PR Description Checklist </summary>
- [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
- [ ] The test plan, such as providing test command.
- [ ] The test results, such as pasting the results comparison before and after, or e2e results
- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model.
- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft.
</details>
**BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)
+89
View File
@@ -0,0 +1,89 @@
#!/bin/bash
# Script to detect changed files and find corresponding test files
# Usage: ./detect_changed_tests.sh [base_sha] [head_sha] [event_type]
#
# Outputs to GITHUB_OUTPUT:
# test_files: Space-separated list of test files to run
# has_tests: 'true' if test files found, 'false' otherwise
set -e
BASE_SHA="${1:-}"
HEAD_SHA="${2:-}"
EVENT_TYPE="${3:-push}"
# Determine base and head SHA based on event type
if [ "$EVENT_TYPE" = "pull_request" ]; then
BASE_SHA="${BASE_SHA:-$GITHUB_BASE_REF}"
HEAD_SHA="${HEAD_SHA:-$GITHUB_SHA}"
elif [ -z "$BASE_SHA" ] || [ -z "$HEAD_SHA" ]; then
# For push events, use git to find the previous commit
BASE_SHA="${BASE_SHA:-$(git rev-parse HEAD~1 2>/dev/null || echo HEAD)}"
HEAD_SHA="${HEAD_SHA:-HEAD}"
fi
echo "Event type: $EVENT_TYPE"
echo "Base SHA: $BASE_SHA"
echo "Head SHA: $HEAD_SHA"
# Get changed Python files
CHANGED_FILES=$(git diff --name-only --diff-filter=ACMRT "$BASE_SHA" "$HEAD_SHA" | grep -E '\.(py)$' || true)
if [ -z "$CHANGED_FILES" ]; then
echo "No Python files changed"
echo "test_files=" >> "$GITHUB_OUTPUT"
echo "has_tests=false" >> "$GITHUB_OUTPUT"
exit 0
fi
echo "Changed Python files:"
echo "$CHANGED_FILES"
# Extract related test files
TEST_FILES=""
for file in $CHANGED_FILES; do
# If changed file is a test file, add it directly
if [[ "$file" == tests/* ]]; then
TEST_FILES="$TEST_FILES $file"
# If changed file is source code, find corresponding test file
elif [[ "$file" == vllm_omni/* ]]; then
# Convert vllm_omni/path/to/module.py to possible test file paths
REL_PATH=$(echo "$file" | sed 's|^vllm_omni/||' | sed 's|\.py$||')
MODULE_NAME=$(basename "$REL_PATH")
DIR_PATH=$(dirname "$REL_PATH")
# Try multiple possible test file naming patterns
POSSIBLE_TESTS=(
"tests/test_${REL_PATH//\//_}.py"
"tests/${REL_PATH//\//_}_test.py"
)
# If multi-level directory, also try directory structure
if [ "$DIR_PATH" != "." ]; then
POSSIBLE_TESTS+=(
"tests/${DIR_PATH}/test_${MODULE_NAME}.py"
"tests/${DIR_PATH}/${MODULE_NAME}_test.py"
)
fi
for test_file in "${POSSIBLE_TESTS[@]}"; do
if [ -f "$test_file" ]; then
TEST_FILES="$TEST_FILES $test_file"
break
fi
done
fi
done
# Deduplicate and format
TEST_FILES=$(echo "$TEST_FILES" | tr ' ' '\n' | grep -v '^$' | sort -u | tr '\n' ' ' | sed 's/^ *//;s/ *$//')
if [ -z "$TEST_FILES" ]; then
echo "No related test files found, will run all tests"
echo "test_files=" >> "$GITHUB_OUTPUT"
echo "has_tests=false" >> "$GITHUB_OUTPUT"
else
echo "Found test files: $TEST_FILES"
echo "test_files=$TEST_FILES" >> "$GITHUB_OUTPUT"
echo "has_tests=true" >> "$GITHUB_OUTPUT"
fi
+61
View File
@@ -0,0 +1,61 @@
name: Build Wheel
on:
push:
branches:
- main
paths-ignore:
- 'docs/**'
- 'examples/**'
- '.github/**'
- '.buildkite/**'
- '**.md'
- '*.yml'
- '.gitignore'
- '.gitattributes'
- 'LICENSE'
pull_request:
branches:
- main
paths-ignore:
- 'docs/**'
- 'examples/**'
- '.github/**'
- '.buildkite/**'
- '*.yml'
- '**.md'
- '.gitignore'
- '.gitattributes'
- 'LICENSE'
workflow_dispatch:
jobs:
build:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ["3.11", "3.12"]
steps:
- uses: actions/checkout@v4
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}
- name: Install uv
run: |
pip install uv
- name: Build Wheel
env:
UV_SYSTEM_PYTHON: 1
run: |
bash scripts/build_wheel.sh --python python
- name: Upload Wheel Artifact
uses: actions/upload-artifact@v4
with:
name: vllm-omni-wheel-py${{ matrix.python-version }}
path: dist/*.whl
+26
View File
@@ -0,0 +1,26 @@
name: pre-commit
on:
pull_request:
branches: [main]
push:
branches: [main]
workflow_dispatch: # for manual trigger
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
permissions:
contents: read
jobs:
pre-commit:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: "3.12"
cache: 'pip'
- uses: pre-commit/action@v3.0.1
+235
View File
@@ -0,0 +1,235 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
vllm/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
.pybuilder/
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
# For a library or package, you might want to ignore these files since the code is
# intended to run in multiple environments; otherwise, check them in:
# .python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# poetry
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
#poetry.lock
# pdm
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
#pdm.lock
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
# in version control.
# https://pdm.fming.dev/#use-with-ide
.pdm.toml
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# pytype static type analyzer
.pytype/
# Cython debug symbols
cython_debug/
# cursor
.cursor/
# scripts
/scripts/
# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be added to the global gitignore or merged into this project gitignore. For a PyCharm
# project, it is recommended to include the following files:
# .idea/
# *.iml
# *.ipr
# *.iws
# .idea/
# vLLM-Omni specific
# Model files and checkpoints
checkpoints/
*.bin
*.safetensors
*.pt
*.pth
# Cache directories
cache/
!vllm_omni/diffusion/cache/
.cache/
diffusion_cache/
kv_cache/
# Logs
logs/
*.log
# Temporary files
tmp/
temp/
.tmp/
# IDE files
.vscode/
.idea/
*.swp
*.swo
*~
# OS files
.DS_Store
.DS_Store?
._*
.Spotlight-V100
.Trashes
ehthumbs.db
Thumbs.db
# Data files
data/
datasets/
*.csv
*.json
*.jsonl
*.parquet
# Output files
outputs/
results/
generated/
output_*/
# Configuration overrides
configs/local.yaml
configs/production.yaml
configs/development.yaml
# Docker
.dockerignore
Dockerfile.dev
discussion
tmp_test
+70
View File
@@ -0,0 +1,70 @@
default_install_hook_types:
- pre-commit
- commit-msg
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v6.0.0
hooks:
# list of supported hooks: https://pre-commit.com/hooks.html
- id: check-yaml
args: ["--unsafe"]
- id: debug-statements
- id: end-of-file-fixer
- id: mixed-line-ending
args: ["--fix=lf"]
- id: trailing-whitespace
args: ["--markdown-linebreak-ext=md"]
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.14.10
hooks:
- id: ruff-check
args: [--output-format, github, --fix]
- id: ruff-format
- repo: https://github.com/crate-ci/typos
rev: typos-dict-v0.13.13
hooks:
- id: typos
# only for staged files
- repo: https://github.com/rhysd/actionlint
# v1.7.8+ sets `go 1.24.0` in go.mod, which older Go toolchains (and most
# current CI images) cannot parse. Pin to v1.7.7 until actionlint fixes the
# go.mod directive.
rev: v1.7.7
hooks:
- id: actionlint
files: ^\.github/workflows/.*\.ya?ml$
- repo: local
hooks:
- id: signoff-commit
name: Sign-off Commit
entry: bash
args:
- -c
- |
if ! grep -q "^Signed-off-by: $(git config user.name) <$(git config user.email)>" "$(git rev-parse --git-path COMMIT_EDITMSG)"; then
printf "\nSigned-off-by: $(git config user.name) <$(git config user.email)>\n" >> "$(git rev-parse --git-path COMMIT_EDITMSG)"
fi
language: system
verbose: true
stages: [commit-msg]
# Keep `suggestion` last
- id: suggestion
name: Suggestion
entry: bash -c 'echo "To bypass all the pre-commit hooks, add --no-verify to git commit. To skip a specific hook, prefix the commit command with SKIP=<hook-id>."'
language: system
verbose: true
pass_filenames: false
# Insert new entries above the `suggestion` entry
- id: check-pickle-imports
name: Prevent new pickle/cloudpickle imports
entry: python tools/pre_commit/check_pickle_imports.py
language: python
types: [python]
additional_dependencies: [regex]
+24
View File
@@ -0,0 +1,24 @@
# Read the Docs configuration file
# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
version: 2
build:
os: ubuntu-22.04
tools:
python: "3.12"
jobs:
post_checkout:
- git fetch --unshallow || true
mkdocs:
configuration: mkdocs.yml
fail_on_warning: true
# Optionally declare the Python requirements required to build your docs
python:
install:
- method: pip
path: .
extra_requirements:
- docs
+3
View File
@@ -0,0 +1,3 @@
# Contributing to vLLM-Omni
You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/)
+201
View File
@@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
+78
View File
@@ -0,0 +1,78 @@
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
<img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
</picture>
</p>
<h3 align="center">
Easy, fast, and cheap omni-modality model serving for everyone
</h3>
<p align="center">
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> |
</p>
---
*Latest News* 🔥
- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm), please check our latest [design](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true).
- [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
---
## About
[vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni is a framework that extends its support for omni-modality model inference and serving:
- **Omni-modality**: Text, image, video, and audio data processing
- **Non-autoregressive Architectures**: extend the AR support of vLLM to Diffusion Transformers (DiT) and other parallel generation models
- **Heterogeneous outputs**: from traditional text generation to multimodal outputs
<p align="center">
<picture>
<img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/omni-modality-model-architecture.png" width=55%>
</picture>
</p>
vLLM-Omni is fast with:
- State-of-the-art AR support by leveraging efficient KV cache management from vLLM
- Pipelined stage execution overlapping for high throughput performance
- Fully disaggregation based on OmniConnector and dynamic resource allocation across stages
vLLM-Omni is flexible and easy to use with:
- Heterogeneous pipeline abstraction to manage complex model workflows
- Seamless integration with popular Hugging Face models
- Tensor, pipeline, data and expert parallelism support for distributed inference
- Streaming outputs
- OpenAI-compatible API server
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
- Omni-modality models (e.g. Qwen-Omni)
- Multi-modality generation models (e.g. Qwen-Image)
## Getting Started
Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn more.
- [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/)
- [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/)
- [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/)
## Contributing
We welcome and value any contributions and collaborations.
Please check out [Contributing to vLLM-Omni](https://vllm-omni.readthedocs.io/en/latest/contributing/) for how to get involved.
## Join the Community
Feel free to ask questions, provide feedbacks and discuss with fellow users of vLLM-Omni in `#sig-omni` slack channel at [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai).
## Star History
[![Star History Chart](https://api.star-history.com/svg?repos=vllm-project/vllm-omni&type=date&legend=top-left)](https://www.star-history.com/#vllm-project/vllm-omni&type=date&legend=top-left)
## License
Apache License 2.0, as found in the [LICENSE](./LICENSE) file.
+42
View File
@@ -0,0 +1,42 @@
# Benchmarks Overview and Architecture
This document explains the benchmark architecture across all benchmark assets in this repo. It describes what we measure, and where to find or plug in new scenarios. Per-task details remain in subfolder READMEs (e.g., `benchmarks/<model>/README.md`).
## Scope and goals
- Establish repeatable latency/throughput measurements for multimodal LLM pipelines.
- Provide both HF Transformers (offline) and vLLM-Omni (multi-stage/pipeline) baselines.
- Make it easy to plug in new datasets and models with minimal changes to the runner scripts.
## Dataset and inputs
- Default example: SeedTTS top-100 prompts (`benchmarks/build_dataset/top100.txt`) via `benchmarks/build_dataset/`.
- Extensible: drop in new prompt files or modality-aligned payloads; keep the expected format for the consuming scripts (e.g., one prompt per line).
- If you add a new dataset, document it under `benchmarks/<model>/README.md` and point scripts to your data path.
## Directory layout
- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
## Reference workflows
- **HF Transformers (offline, single process)**
Script (example): `benchmarks/<model>/transformers/eval_qwen3_moe_omni_transformers.sh`
Outputs: `benchmark_results/perf_stats.json`, `benchmark_results/results.json`, `benchmark_results/audio/` (if audio is produced).
- **vLLM-Omni end-to-end pipeline**
Script (example): `benchmarks/<model>/vllm_omni/eval_qwen3_moe_omni.sh`
Outputs: `vllm_omni/logs/*.stats.jsonl` (per-stage/overall latency & TPS), `vllm_omni/logs/stage*.log`, `vllm_omni/outputs/` (text/audio artifacts).
- **Adding a new task/model**
1) Create `benchmarks/<model>/transformers/` and/or `benchmarks/<model>/vllm_omni/` with scripts referencing your model and dataset.
2) Add a task README describing dataset, configs, and expected outputs.
3) Keep the output/log structure similar for easy comparison (perf_stats/results/audio or text outputs; stats.jsonl/logs for pipeline).
## Metrics to watch
- **Throughput**: `overall_tps`, `*_tps_avg` per stage.
- **Latency distribution**: look for long tails in `*.stats.jsonl`.
- **Quality/completeness**: missing outputs or errors in stage logs indicate pipeline failures or misconfigurations.
## Troubleshooting
- Verify GPU/driver/FlashAttention2 requirements for your chosen model/config.
- Ensure network access for dataset/model downloads (Google Drive, Hugging Face, etc.).
- If outputs are missing or slow, inspect per-stage logs and `*.stats.jsonl` for errors, stragglers, or contention.
@@ -0,0 +1,82 @@
# Benchmark Dataset Preparation Guide
This guide describes how to download and prepare the SeedTTS test dataset for benchmarking Qwen-Omni models.
## Prerequisites
- Python 3.8+
- `gdown` for downloading from Google Drive
- Access to the benchmark scripts
## Steps
### 1. Navigate to the Dataset Directory
```bash
cd benchmarks/build_dataset
```
### 2. Install Dependencies
```bash
pip install gdown
```
### 3. Download the SeedTTS Test Dataset
Download the dataset from Google Drive:
```bash
gdown --id 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
```
### 4. Extract the Dataset
```bash
tar -xf seedtts_testset.tar
```
### 5. Prepare the Metadata File
Copy the English metadata file to the working directory:
```bash
cp seedtts_testset/en/meta.lst meta.lst
```
### 6. Extract Prompts
Extract the first N prompts from the metadata file:
```bash
# Extract top 100 prompts (adjust -n for different amounts)
python extract_tts_prompts.py -i meta.lst -o top100.txt -n 100
```
**Options:**
- `-i, --input`: Input metadata file (default: `meta.lst`)
- `-o, --output`: Output prompts file (default: `prompts.txt`)
- `-n, --num_lines`: Number of prompts to extract (required)
### 7. Clean Up (Optional)
Remove temporary files to save disk space:
```bash
rm -rf seedtts_testset
rm seedtts_testset.tar
rm meta.lst
```
## Quick Start (All-in-One)
```bash
# Full setup and benchmark
cd benchmarks/build_dataset
pip install gdown
gdown --id 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
tar -xf seedtts_testset.tar
cp seedtts_testset/en/meta.lst meta.lst
python extract_tts_prompts.py -i meta.lst -o top100.txt -n 100
rm -rf seedtts_testset seedtts_testset.tar meta.lst
```
@@ -0,0 +1,73 @@
#!/usr/bin/env python3
"""
Extract prompts from meta.lst and save them to a txt file.
Each line in meta.lst has the format:
ID|prompt_text|audio_path|target_text
This script extracts the prompt_text (second field) from the first N lines.
"""
import argparse
from pathlib import Path
def extract_prompts(input_file: str, output_file: str, num_lines: int) -> None:
"""
Extract prompts from meta.lst and save to output file.
Args:
input_file: Path to the meta.lst file
output_file: Path to the output txt file
num_lines: Number of lines to process
"""
prompts = []
with open(input_file, encoding="utf-8") as f:
for i, line in enumerate(f):
if i >= num_lines:
break
line = line.strip()
if not line: # Skip empty lines
continue
parts = line.split("|")
if len(parts) >= 2:
prompt = parts[1] # The prompt is the second field
prompts.append(prompt)
# Write prompts to output file
with open(output_file, "w", encoding="utf-8") as f:
for prompt in prompts:
f.write(prompt + "\n")
# Print result stats
print(f"Extracted {len(prompts)} prompts from first {num_lines} lines")
print(f"Saved to: {output_file}")
def main():
parser = argparse.ArgumentParser(description="Extract prompts from meta.lst file")
parser.add_argument(
"-i", "--input", type=str, default="meta.lst", help="Input meta.lst file path (default: meta.lst)"
)
parser.add_argument(
"-o", "--output", type=str, default="prompts.txt", help="Output txt file path (default: prompts.txt)"
)
parser.add_argument(
"-n", "--num_lines", type=int, required=True, help="Number of lines to extract from the beginning"
)
args = parser.parse_args()
# Check if input file exists
if not Path(args.input).exists():
print(f"Error: Input file '{args.input}' not found")
return
extract_prompts(args.input, args.output, args.num_lines)
if __name__ == "__main__":
main()
+117
View File
@@ -0,0 +1,117 @@
# Diffusion Serving Benchmark (Image/Video)
This folder contains an online-serving benchmark script for diffusion models.
It sends requests to a vLLM OpenAI-compatible endpoint and reports throughput,
latency percentiles, and optional SLO attainment.
The main entrypoint is:
- `benchmarks/diffusion/diffusion_benchmark_serving.py`
## 1. Quick Start
1. Start the server:
```bash
vllm serve Qwen/Qwen-Image --omni --port 8099
```
2. Run a minimal benchmark:
```bash
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--base-url http://localhost:8099 \
--model Qwen/Qwen-Image \
--task t2i \
--dataset vbench \
--num-prompts 5
```
**Notes**
- The benchmark talks to `http://<host>:<port>/v1/chat/completions`.
- If you run the server on another host or port, pass `--base-url` accordingly.
## 2. Supported Datasets
The benchmark supports three dataset modes via `--dataset`:
- `vbench`: Built-in prompt/data loader.
- `trace`: Heterogeneous request traces (each request can have different resolution/frames/steps).
- `random`: Synthetic prompts for quick smoke tests.
### VBench dataset
If you use i2v/i2i bench datasets and need auto-download support, you may need:
```bash
uv pip install gdown
```
### Trace dataset
Use `--dataset trace` to replay a trace file. The trace can specify per-request fields such as:
- `width`, `height`
- `num_frames` (video)
- `num_inference_steps`
- `seed`, `fps`
- optional `slo_ms` (per-request SLO target)
By default (when `--dataset-path` is not provided), the script downloads a default trace from
the HuggingFace dataset repo `asukaqaqzz/Dit_Trace`. The default filename can depend on `--task`
(e.g., `t2v` uses a video trace).
Current defaults:
- `--task t2i` -> `sd3_trace.txt`
- `--task t2v` -> `cogvideox_trace.txt`
You can point to your own trace using `--dataset-path`.
## 3. Benchmark Parameters
### Basic flags
- `--base-url`: Server address (the script calls `.../v1/chat/completions`).
- `--model`: The OpenAI-compatible `model` field.
- `--task`: Task type (e.g., `t2i`, `t2v`, `i2i`, `i2v`).
- `--dataset`: Dataset mode (`vbench` / `trace` / `random`).
- `--num-prompts`: Number of requests to send.
Common optional flags:
- `--output-file`: Write metrics to a JSON file.
- `--disable-tqdm`: Disable the progress bar.
### Resolution / frames / steps: CLI defaults vs dataset fields
Related flags: `--width`, `--height`, `--num-frames`, `--fps`, `--num-inference-steps`.
- For `vbench` / `random`: these CLI flags act as global defaults for all generated requests.
- For `trace`: each request can carry its own fields (e.g., `width/height/num_frames/num_inference_steps`).
Precedence rules for `trace` (i.e., what actually gets sent):
- `width/height`: if either `--width` or `--height` is explicitly set, it overrides per-request values from the trace; otherwise per-request values are used when present.
- `num_frames`: per-request `num_frames` takes precedence; otherwise fall back to `--num-frames`.
- `num_inference_steps`: per-request `num_inference_steps` takes precedence; otherwise fall back to `--num-inference-steps`.
### SLO, warmup, and max concurrency
Enable SLO evaluation with `--slo`.
- If a request in the trace already has `slo_ms`, that value is used.
- Otherwise, the script runs warmup requests to infer a base unit time, estimates `expected_ms` by linearly scaling with area/frames/steps, and then sets `slo_ms = expected_ms * --slo-scale`.
Warmup flags:
- `--warmup-requests`: Number of warmup requests.
- `--warmup-num-inference-steps`: Steps used during warmup.
- For `--task t2v`: warmup requests are forced to use `num_frames=1` to make warmup faster and less noisy.
Traffic / concurrency flags:
- `--request-rate`: Target request rate (requests/second). If set to `inf`, the script sends all requests immediately.
- `--max-concurrency`: Max number of in-flight requests (default: `1`). This can hard-cap the achieved QPS: if it is too small, requests will queue behind the semaphore, and both achieved throughput and observed SLO attainment can be skewed.
File diff suppressed because it is too large. Load diff
+86
View File
@@ -0,0 +1,86 @@
# Benchmarks Guide
This README explains how to (1) prepare benchmark datasets and (2) run the provided Qwen3-Omni benchmarks.
## 1) Prepare the dataset (SeedTTS top100)
```bash
cd benchmarks/build_dataset
pip install gdown
# Download SeedTTS test set from Google Drive
gdown --id 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
# Extract
tar -xf seedtts_testset.tar
# Copy metadata and extract top-100 prompts
cp seedtts_testset/en/meta.lst meta.lst
python extract_prompts.py -i meta.lst -o top100.txt -n 100
# (Optional) clean up to save space
rm -rf seedtts_testset seedtts_testset.tar meta.lst
```
Artifacts:
- `benchmarks/build_dataset/top100.txt` — 100 text prompts (one per line).
## 2) Run benchmarks
All commands assume repo root (`vllm-omni`).
### A. Transformers benchmark (offline, HF Transformers)
```
bash benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh
```
What it does:
- Runs `qwen3_omni_moe_transformers.py` over `top100.txt` with `--num_prompts 100`.
- Outputs to `benchmarks/qwen3-omni/transformers/benchmark_results/`:
- `perf_stats.json` — aggregated & per-prompt TPS/latency (thinker/talker/code2wav/overall).
- `results.json` — per-prompt outputs and audio paths.
- `audio/` — ~100 generated `.wav` files.
Key checks:
- `overall_tps` and `*_tps_avg` should be non-zero and reasonably stable.
- Investigate any 0/NaN or unusually low TPS / long-tail latency.
### B. vLLM Omni end-to-end benchmark (pipeline)
```
bash benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh
```
What it does:
- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--enable-stats`.
- Uses `benchmarks/build_dataset/top100.txt` and writes to:
- Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
- `omni_llm_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
- `omni_llm_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
- `omni_llm_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
- Outputs: `benchmarks/qwen3-omni/vllm_omni/outputs/` — ~100 text and `.wav` files.
Key checks:
- Overall stats: end-to-end latency/TPS should be reasonable.
- Orchestrator stats: per-stage latency should be stable; investigate long tails.
- Stage logs: ensure no errors and no unusually slow stages.
## Performance snapshot
The chart below summarizes our measured Qwen3-Omni MoE end-to-end benchmark, comparing vLLM-Omni against HF Transformers. It shows the overall throughput advantage for vLLM-Omni. These are actual experiment results—please refer to this performance when evaluating or reproducing the benchmark.
![vLLM-Omni vs HF](./vllm-omni-vs-hf.png)
## Directory layout
- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
- `benchmarks/<model>/vllm-omni-vs-hf.png` — current performance snapshot (overall throughput comparison).
- `benchmarks/<model>/transformers/` — HF Transformers benchmarks (offline reference).
## Troubleshooting
- Make sure GPU/driver/FlashAttention2 requirements are met for the chosen model.
- If downloads fail, confirm network access to Google Drive (`gdown`) and Hugging Face.
- If audio files are missing, check for errors in stage logs or model generation.***
@@ -0,0 +1,26 @@
#!/bin/bash
# Qwen3-Omni Transformers Benchmark Evaluation Script
# This script must be run from the vllm-omni root directory
# Get the directory where this script is located
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# Navigate to vllm-omni root directory (4 levels up from script location)
VLLM_OMNI_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
cd "$VLLM_OMNI_ROOT" || { echo "Error: Failed to navigate to vllm-omni directory"; exit 1; }
echo "Working directory: $(pwd)"
# Verify we're in the correct directory and run benchmark
if [[ ! -f "benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py" ]]; then
echo "Error: Not in vllm-omni root directory. Please run from vllm-omni folder."
else
cd benchmarks/qwen3-omni/transformers
python qwen3_omni_moe_transformers.py --prompts_file ../../build_dataset/top100.txt --num_prompts 100
echo "Logs and outputs are saved to $(pwd)/benchmark_results:"
echo " - perf_stats.json Aggregated/per-prompt TPS and latency (thinker/talker/code2wav/overall)"
echo " - results.json Per-prompt outputs and audio paths"
echo " - audio/ Generated wav files, there should be 100 wav file generated"
echo "Key checks: overall_tps and *_tps_avg should be non-zero and stable; investigate 0/NaN or unusually low TPS/long-tail latency."
fi
@@ -0,0 +1,265 @@
import time
import torch
from transformers import Qwen3OmniMoeForConditionalGeneration
class Qwen3OmniMoeForConditionalGenerationWithLogging(Qwen3OmniMoeForConditionalGeneration):
@torch.no_grad()
def generate(
self,
input_ids: torch.Tensor | None = None,
speaker: str = "Ethan",
use_audio_in_video: bool = False,
return_audio: bool | None = None,
thinker_max_new_tokens: int = 1024,
thinker_eos_token_id: int = 151645,
talker_max_new_tokens: int = 4096,
talker_do_sample: bool = True,
talker_top_k: int = 50,
talker_top_p: float = 1.0,
talker_temperature: float = 0.9,
talker_repetition_penalty: float = 1.05,
**kwargs,
):
total_t0 = time.time()
perf_stats = {
"thinker_tokens": 0,
"thinker_time_s": 0.0,
"thinker_tps": 0.0,
"talker_tokens": 0,
"talker_time_s": 0.0,
"talker_tps": 0.0,
"code2wav_tokens": 0,
"code2wav_time_s": 0.0,
"code2wav_tps": 0.0,
"total_tokens": 0,
"total_time_s": 0.0,
"total_tps": 0.0,
}
if return_audio and not self.has_talker:
raise ValueError(
"Cannot use talker when talker module not initialized. "
"Use `enable_talker` method or set enable_talker in config "
"to enable talker."
)
if return_audio is None:
return_audio = self.has_talker
shared_kwargs = {"use_audio_in_video": use_audio_in_video}
thinker_kwargs = {
"max_new_tokens": thinker_max_new_tokens,
"eos_token_id": thinker_eos_token_id,
}
talker_kwargs = {}
token2wav_kwargs = {}
if return_audio:
speaker_id = self.config.talker_config.speaker_id.get(speaker.lower())
if speaker_id is None:
raise NotImplementedError(f"Speaker {speaker} not implemented")
if input_ids.shape[0] != 1:
raise NotImplementedError("Qwen3-Omni currently does not support batched inference with audio output")
talker_suppressed_tokens = [
i
for i in range(
self.config.talker_config.text_config.vocab_size - 1024,
self.config.talker_config.text_config.vocab_size,
)
if i != self.config.talker_config.codec_eos_token_id
] # Suppress additional special tokens, should not be predicted
talker_kwargs = {
"max_new_tokens": talker_max_new_tokens,
"do_sample": talker_do_sample,
"top_k": talker_top_k,
"top_p": talker_top_p,
"temperature": talker_temperature,
"eos_token_id": self.config.talker_config.codec_eos_token_id,
"repetition_penalty": talker_repetition_penalty,
"suppress_tokens": talker_suppressed_tokens,
"output_hidden_states": True,
"return_dict_in_generate": True,
}
token2wav_kwargs = {}
for key, value in kwargs.items():
if key.startswith("thinker_"):
thinker_kwargs[key[len("thinker_") :]] = value
elif key.startswith("talker_"):
talker_kwargs[key[len("talker_") :]] = value
elif key.startswith("token2wav_"):
token2wav_kwargs[key[len("token2wav_") :]] = value
# Process special input values
elif key == "feature_attention_mask":
thinker_kwargs[key] = value
talker_kwargs["audio_feature_lengths"] = torch.sum(value, dim=1)
elif key in ("input_features", "attention_mask"):
thinker_kwargs[key] = value
# Put other key to shared kwargs
else:
shared_kwargs[key] = value
# Merge kwargs
for key, value in shared_kwargs.items():
if key not in thinker_kwargs:
thinker_kwargs[key] = value
if key not in talker_kwargs and key in ["image_grid_thw", "video_grid_thw", "video_second_per_grid"]:
talker_kwargs[key] = value
if key not in token2wav_kwargs:
token2wav_kwargs[key] = value
# 1. Generate from thinker module
generate_audio = return_audio and self.has_talker
if generate_audio:
thinker_kwargs["output_hidden_states"] = True
thinker_kwargs["return_dict_in_generate"] = True
t0 = time.time()
thinker_result = self.thinker.generate(input_ids=input_ids, **thinker_kwargs)
t1 = time.time()
perf_stats["thinker_time_s"] = max(0.0, t1 - t0)
try:
prompt_len = int(input_ids.shape[1]) if input_ids is not None else 0
total_len = int(thinker_result.sequences.shape[-1])
thinker_out_len = max(0, total_len - prompt_len)
except Exception:
thinker_out_len = 0
perf_stats["thinker_tokens"] = thinker_out_len
perf_stats["thinker_tps"] = (
(thinker_out_len / perf_stats["thinker_time_s"]) if perf_stats["thinker_time_s"] > 0 else 0.0
)
if not generate_audio:
perf_stats["total_tokens"] = perf_stats["thinker_tokens"]
perf_stats["total_time_s"] = time.time() - total_t0
perf_stats["total_tps"] = (
(perf_stats["total_tokens"] / perf_stats["total_time_s"]) if perf_stats["total_time_s"] > 0 else 0.0
)
# attach stats to self
setattr(self, "_perf_stats_last", perf_stats)
if not hasattr(self, "_perf_stats_history"):
setattr(self, "_perf_stats_history", [])
self._perf_stats_history.append(perf_stats)
return thinker_result, None
# 2. Prepare talker input
thinker_embed = torch.cat([hidden_states[0] for hidden_states in thinker_result.hidden_states], dim=1).to(
self.talker.device
) # [1 t d]
thinker_hidden = torch.cat(
[
hidden_states[self.config.talker_config.accept_hidden_layer]
for hidden_states in thinker_result.hidden_states
],
dim=1,
).to(self.talker.device) # [1 t d]
im_start_indexes = torch.cat(
(
torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(),
torch.tensor([thinker_result.sequences.shape[-1]], device=input_ids.device, dtype=input_ids.dtype),
),
dim=-1,
).to(self.talker.device) # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here.
multimodal_mask = (
(thinker_result.sequences == self.config.thinker_config.audio_token_id) |
(thinker_result.sequences == self.config.thinker_config.image_token_id) |
(thinker_result.sequences == self.config.thinker_config.video_token_id)
).to(self.talker.device) # [1 t] # fmt: skip
talker_special_tokens = torch.tensor(
[[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]],
device=self.thinker.device,
dtype=input_ids.dtype,
)
tts_bos_embed, tts_eos_embed, tts_pad_embed = (
self.talker.text_projection(self.thinker.get_input_embeddings()(talker_special_tokens))
.to(self.talker.device)
.chunk(3, dim=1)
) # 3 * [1 1 d]
talker_input_embeds = [] # [1 t d]
talker_input_ids = []
# For every chatml parts
for i in range(len(im_start_indexes) - 1):
im_start_index = im_start_indexes[i]
segment_end_index = im_start_indexes[i + 1]
role_token = input_ids[0][im_start_index + 1]
# Talker should ignore thinker system prompt
if role_token == self.config.system_token_id:
continue
# Talker takes word embeddings for tokens and hidden state from `accept_hidden_layer` for multimodal inputs
elif role_token == self.config.user_token_id:
talker_user_part = self._get_talker_user_parts(
im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed
)
talker_input_embeds.append(talker_user_part)
talker_input_ids.append(thinker_result.sequences[:, im_start_index:segment_end_index])
# Take assistant output (for now)
elif role_token == self.config.assistant_token_id and i == len(im_start_indexes) - 2:
talker_assistant_embeds, talker_assistant_ids, trailing_text_hidden = self._get_talker_assistant_parts(
im_start_index,
segment_end_index,
speaker_id,
thinker_embed,
tts_pad_embed,
tts_bos_embed,
tts_eos_embed,
)
talker_input_embeds.append(talker_assistant_embeds)
talker_input_ids.append(talker_assistant_ids)
# History assistant output (ignore for now)
elif role_token == self.config.assistant_token_id and i != len(im_start_indexes) - 2:
continue
else:
raise AssertionError("Expect role id after <|im_start|> (assistant, user, system)")
talker_input_embed = torch.cat([embed.to(self.talker.device) for embed in talker_input_embeds], dim=1)
talker_input_id = torch.cat([embed.to(self.talker.device) for embed in talker_input_ids], dim=1)
t2 = time.time()
talker_result = self.talker.generate(
inputs_embeds=talker_input_embed,
trailing_text_hidden=trailing_text_hidden,
tts_pad_embed=tts_pad_embed,
talker_input_ids=talker_input_id, # Not use input_ids to prevent repetition penalty out of bound
**talker_kwargs,
)
t3 = time.time()
perf_stats["talker_time_s"] = max(0.0, t3 - t2)
talker_codes = (
torch.stack([hid[-1] for hid in talker_result.hidden_states if hid[-1] is not None], dim=1)
.transpose(1, 2)
.to(self.code2wav.device)
)
try:
# codes shape: (B, num_quantizers, T). We log T as token length.
perf_stats["talker_tokens"] = int(talker_codes.shape[-1])
except Exception:
perf_stats["talker_tokens"] = 0
perf_stats["talker_tps"] = (
(perf_stats["talker_tokens"] / perf_stats["talker_time_s"]) if perf_stats["talker_time_s"] > 0 else 0.0
)
t4 = time.time()
talker_wavs = self.code2wav.chunked_decode(talker_codes, chunk_size=300, left_context_size=25).float()
t5 = time.time()
perf_stats["code2wav_time_s"] = max(0.0, t5 - t4)
perf_stats["code2wav_tokens"] = perf_stats["talker_tokens"] # same T, not times 16
perf_stats["code2wav_tps"] = (
(perf_stats["code2wav_tokens"] / perf_stats["code2wav_time_s"])
if perf_stats["code2wav_time_s"] > 0
else 0.0
)
perf_stats["total_tokens"] = perf_stats["thinker_tokens"] + perf_stats["talker_tokens"]
perf_stats["total_time_s"] = time.time() - total_t0
perf_stats["total_tps"] = (
(perf_stats["total_tokens"] / perf_stats["total_time_s"]) if perf_stats["total_time_s"] > 0 else 0.0
)
setattr(self, "_perf_stats_last", perf_stats)
if not hasattr(self, "_perf_stats_history"):
setattr(self, "_perf_stats_history", [])
self._perf_stats_history.append(perf_stats)
return thinker_result, talker_wavs.float()
__all__ = [
"Qwen3OmniMoeForConditionalGenerationWithLogging",
]
@@ -0,0 +1,275 @@
import argparse
import json
import os
import soundfile as sf
from qwen3_omni_moe_model import Qwen3OmniMoeForConditionalGenerationWithLogging
from qwen_omni_utils import process_mm_info
from tqdm import tqdm
from transformers import Qwen3OmniMoeProcessor
MODEL_PATH = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
# MODEL_PATH = "Qwen/Qwen3-Omni-30B-A3B-Thinking"
def load_prompts(prompts_file: str) -> list[str]:
"""Load prompts from a text file, one prompt per line."""
prompts = []
with open(prompts_file, encoding="utf-8") as f:
for line in f:
line = line.strip()
if line:
prompts.append(line)
return prompts
def run_benchmark(
model,
processor,
prompts: list[str],
output_dir: str = "benchmark_results",
speaker: str = "Ethan",
use_audio_in_video: bool = True,
):
"""
Run benchmark on a list of prompts and collect performance stats.
Args:
model: The Qwen3OmniMoe model
processor: The Qwen3OmniMoe processor
prompts: List of text prompts to process
output_dir: Directory to save results
speaker: Speaker voice for audio output
use_audio_in_video: Whether to use audio in video
Returns:
tuple: (aggregated_stats, results, audio_outputs)
- aggregated_stats: dict with aggregated performance statistics
- results: list of dicts with per-prompt results
- audio_outputs: list of audio tensors/arrays (or None if no audio)
"""
os.makedirs(output_dir, exist_ok=True)
audio_dir = os.path.join(output_dir, "audio")
os.makedirs(audio_dir, exist_ok=True)
all_stats = []
results = []
audio_outputs = []
for idx, prompt in enumerate(tqdm(prompts, desc="Processing prompts")):
conversation = [
{
"role": "user",
"content": [{"type": "text", "text": prompt}],
},
]
# Preparation for inference
text = processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
audios, images, videos = process_mm_info(conversation, use_audio_in_video=use_audio_in_video)
inputs = processor(
text=text,
audio=audios,
images=images,
videos=videos,
return_tensors="pt",
padding=True,
use_audio_in_video=use_audio_in_video,
)
inputs = inputs.to(model.device).to(model.dtype)
# Inference: Generation of the output text and audio
text_ids, audio = model.generate(
**inputs, speaker=speaker, thinker_return_dict_in_generate=True, use_audio_in_video=use_audio_in_video
)
# Decode output text
output_text = processor.batch_decode(
text_ids.sequences[:, inputs["input_ids"].shape[1] :],
skip_special_tokens=True,
clean_up_tokenization_spaces=False,
)[0]
# Collect performance stats
perf_stats = None
if hasattr(model, "_perf_stats_last"):
perf_stats = model._perf_stats_last.copy()
perf_stats["prompt_idx"] = idx
perf_stats["prompt"] = prompt
all_stats.append(perf_stats)
# Save audio and collect audio output
audio_path = None
audio_data = None
if audio is not None:
audio_data = audio.reshape(-1).detach().cpu().numpy()
audio_path = os.path.join(audio_dir, f"output_{idx:04d}.wav")
sf.write(
audio_path,
audio_data,
samplerate=24000,
)
audio_outputs.append(audio_data)
else:
audio_outputs.append(None)
# Save result
result = {
"idx": idx,
"prompt": prompt,
"output": output_text,
"audio_path": audio_path,
"perf_stats": perf_stats,
}
results.append(result)
# Aggregate statistics
aggregated_stats = aggregate_stats(all_stats)
# Save all results
results_path = os.path.join(output_dir, "results.json")
with open(results_path, "w", encoding="utf-8") as f:
json.dump(results, f, ensure_ascii=False, indent=2)
# Save aggregated stats
stats_path = os.path.join(output_dir, "perf_stats.json")
with open(stats_path, "w", encoding="utf-8") as f:
json.dump({"aggregated": aggregated_stats, "per_prompt": all_stats}, f, ensure_ascii=False, indent=2)
# Count saved audio files
num_audio_saved = sum(1 for a in audio_outputs if a is not None)
print(f"\nSaved {num_audio_saved} audio files to {audio_dir}/")
return aggregated_stats, results, audio_outputs
def aggregate_stats(all_stats: list[dict]) -> dict:
"""Aggregate performance statistics from multiple runs."""
if not all_stats:
return {}
keys = [
"thinker_tokens",
"thinker_time_s",
"thinker_tps",
"talker_tokens",
"talker_time_s",
"talker_tps",
"code2wav_tokens",
"code2wav_time_s",
"code2wav_tps",
"total_tokens",
"total_time_s",
"total_tps",
]
aggregated = {
"num_samples": len(all_stats),
}
for key in keys:
values = [s.get(key, 0) for s in all_stats if key in s]
if values:
aggregated[f"{key}_sum"] = sum(values)
aggregated[f"{key}_avg"] = sum(values) / len(values)
aggregated[f"{key}_min"] = min(values)
aggregated[f"{key}_max"] = max(values)
# Calculate overall throughput
total_tokens = aggregated.get("total_tokens_sum", 0)
total_time = aggregated.get("total_time_s_sum", 0)
if total_time > 0:
aggregated["overall_tps"] = total_tokens / total_time
return aggregated
def print_stats(stats: dict):
"""Print performance statistics in a formatted way."""
print("\n" + "=" * 60)
print("Performance Statistics Summary")
print("=" * 60)
print(f"\nNumber of samples: {stats.get('num_samples', 0)}")
print("\n--- Thinker ---")
print(f" Total tokens: {stats.get('thinker_tokens_sum', 0):.0f}")
print(f" Total time: {stats.get('thinker_time_s_sum', 0):.2f}s")
print(f" Avg TPS: {stats.get('thinker_tps_avg', 0):.2f}")
print(f" Min TPS: {stats.get('thinker_tps_min', 0):.2f}")
print(f" Max TPS: {stats.get('thinker_tps_max', 0):.2f}")
print("\n--- Talker ---")
print(f" Total tokens: {stats.get('talker_tokens_sum', 0):.0f}")
print(f" Total time: {stats.get('talker_time_s_sum', 0):.2f}s")
print(f" Avg TPS: {stats.get('talker_tps_avg', 0):.2f}")
print(f" Min TPS: {stats.get('talker_tps_min', 0):.2f}")
print(f" Max TPS: {stats.get('talker_tps_max', 0):.2f}")
print("\n--- Code2Wav ---")
print(f" Total tokens: {stats.get('code2wav_tokens_sum', 0):.0f}")
print(f" Total time: {stats.get('code2wav_time_s_sum', 0):.2f}s")
print(f" Avg TPS: {stats.get('code2wav_tps_avg', 0):.2f}")
print(f" Min TPS: {stats.get('code2wav_tps_min', 0):.2f}")
print(f" Max TPS: {stats.get('code2wav_tps_max', 0):.2f}")
print("\n--- Overall ---")
print(f" Total tokens: {stats.get('total_tokens_sum', 0):.0f}")
print(f" Total time: {stats.get('total_time_s_sum', 0):.2f}s")
print(f" Overall TPS: {stats.get('overall_tps', 0):.2f}")
print(f" Avg TPS: {stats.get('total_tps_avg', 0):.2f}")
print(f" Min TPS: {stats.get('total_tps_min', 0):.2f}")
print(f" Max TPS: {stats.get('total_tps_max', 0):.2f}")
print("=" * 60 + "\n")
def main():
parser = argparse.ArgumentParser(description="Qwen3-Omni Benchmark Script")
parser.add_argument(
"--prompts_file",
type=str,
default="benchmark/build_dataset/top100.txt",
help="Path to the prompts file (one prompt per line)",
)
parser.add_argument(
"--output_dir", type=str, default="benchmark_results", help="Directory to save benchmark results"
)
parser.add_argument("--model_path", type=str, default=MODEL_PATH, help="Path to the model")
parser.add_argument("--speaker", type=str, default="Ethan", help="Speaker voice for audio output")
parser.add_argument("--num_prompts", type=int, default=None, help="Number of prompts to process (default: all)")
args = parser.parse_args()
# Load model and processor
print(f"Loading model from {args.model_path}...")
model = Qwen3OmniMoeForConditionalGenerationWithLogging.from_pretrained(
args.model_path,
dtype="auto",
device_map="auto",
attn_implementation="flash_attention_2",
)
processor = Qwen3OmniMoeProcessor.from_pretrained(args.model_path)
# Benchmark mode
print(f"Loading prompts from {args.prompts_file}...")
prompts = load_prompts(args.prompts_file)
if args.num_prompts:
prompts = prompts[: args.num_prompts]
print(f"Running benchmark on {len(prompts)} prompts...")
aggregated_stats, results, audio_outputs = run_benchmark(
model=model,
processor=processor,
prompts=prompts,
output_dir=args.output_dir,
speaker=args.speaker,
)
print_stats(aggregated_stats)
print(f"\nResults saved to {args.output_dir}/")
if __name__ == "__main__":
main()
Binary file not shown.

After

Width:  |  Height:  |  Size: 115 KiB

@@ -0,0 +1,37 @@
#!/bin/bash
# Qwen3-Omni Benchmark Evaluation Script
# This script must be run from the vllm-omni root directory
# Get the directory where this script is located
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# Navigate to vllm-omni root directory (4 levels up from script location)
VLLM_OMNI_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
cd "$VLLM_OMNI_ROOT" || { echo "Error: Failed to navigate to vllm-omni directory"; exit 1; }
echo "Working directory: $(pwd)"
# Verify we're in the correct directory and run benchmark
if [[ ! -d "benchmarks/qwen3-omni/vllm_omni" ]]; then
echo "Error: Not in vllm-omni root directory. Please run from vllm-omni folder."
else
log_dir=benchmarks/qwen3-omni/vllm_omni/logs
outputs_dir=benchmarks/qwen3-omni/vllm_omni/outputs
end2end_script_path=examples/offline_inference/qwen3_omni/end2end.py
build_dataset_path=benchmarks/build_dataset/top100.txt
python $end2end_script_path --output-wav $outputs_dir \
--query-type text \
--txt-prompts $build_dataset_path \
--enable-stats \
--log-dir $log_dir
echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
echo " - omni_llm_pipeline_text run dir/base name"
echo " - omni_llm_pipeline_text.orchestrator.stats.jsonl orchestrator-stage latency stats"
echo " - omni_llm_pipeline_text.overall.stats.jsonl overall latency/TPS stats"
echo " - omni_llm_pipeline_text.stage0.log per-stage detailed logs"
echo " - omni_llm_pipeline_text.stage1.log"
echo " - omni_llm_pipeline_text.stage2.log"
echo "Key checks: overall.stats.jsonl for end-to-end latency/TPS; orchestrator.stats.jsonl for stable per-stage latency; stage*.log for errors or long tails."
echo " - outputs/ Generated txt and wav files, there should be 100 text and wav files generated respectively"
fi
+760
View File
@@ -0,0 +1,760 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# ruff: noqa
# code borrowed from https://github.com/pytorch/pytorch/blob/main/torch/utils/collect_env.py
import datetime
import locale
import os
import subprocess
import sys
# Unlike the rest of the PyTorch this file must be python2 compliant.
# This script outputs relevant system environment info
# Run it with `python collect_env.py` or `python -m torch.utils.collect_env`
from collections import namedtuple
import regex as re
from vllm.envs import environment_variables
try:
import torch
TORCH_AVAILABLE = True
except (ImportError, NameError, AttributeError, OSError):
TORCH_AVAILABLE = False
# System Environment Information
SystemEnv = namedtuple(
"SystemEnv",
[
"torch_version",
"is_debug_build",
"cuda_compiled_version",
"gcc_version",
"clang_version",
"cmake_version",
"os",
"libc_version",
"python_version",
"python_platform",
"is_cuda_available",
"cuda_runtime_version",
"cuda_module_loading",
"nvidia_driver_version",
"nvidia_gpu_models",
"cudnn_version",
"pip_version", # 'pip' or 'pip3'
"pip_packages",
"conda_packages",
"hip_compiled_version",
"hip_runtime_version",
"miopen_runtime_version",
"caching_allocator_config",
"is_xnnpack_available",
"cpu_info",
"rocm_version", # vllm specific field
"vllm_version", # vllm specific field
"vllm_omni_version", # vllm-omni specific field
"vllm_build_flags", # vllm specific field
"gpu_topo", # vllm specific field
"env_vars",
],
)
DEFAULT_CONDA_PATTERNS = {
"torch",
"numpy",
"cudatoolkit",
"soumith",
"mkl",
"magma",
"triton",
"optree",
"nccl",
"transformers",
"zmq",
"nvidia",
"pynvml",
"flashinfer-python",
}
DEFAULT_PIP_PATTERNS = {
"torch",
"numpy",
"mypy",
"flake8",
"triton",
"optree",
"onnx",
"nccl",
"transformers",
"zmq",
"nvidia",
"pynvml",
"flashinfer-python",
}
def run(command):
"""Return (return-code, stdout, stderr)."""
shell = True if type(command) is str else False
try:
p = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=shell)
raw_output, raw_err = p.communicate()
rc = p.returncode
if get_platform() == "win32":
enc = "oem"
else:
enc = locale.getpreferredencoding()
output = raw_output.decode(enc)
if command == "nvidia-smi topo -m":
# don't remove the leading whitespace of `nvidia-smi topo -m`
# because they are meaningful
output = output.rstrip()
else:
output = output.strip()
err = raw_err.decode(enc)
return rc, output, err.strip()
except FileNotFoundError:
cmd_str = command if isinstance(command, str) else command[0]
return 127, "", f"Command not found: {cmd_str}"
def run_and_read_all(run_lambda, command):
"""Run command using run_lambda; reads and returns entire output if rc is 0."""
rc, out, _ = run_lambda(command)
if rc != 0:
return None
return out
def run_and_parse_first_match(run_lambda, command, regex):
"""Run command using run_lambda, returns the first regex match if it exists."""
rc, out, _ = run_lambda(command)
if rc != 0:
return None
match = re.search(regex, out)
if match is None:
return None
return match.group(1)
def run_and_return_first_line(run_lambda, command):
"""Run command using run_lambda and returns first line if output is not empty."""
rc, out, _ = run_lambda(command)
if rc != 0:
return None
return out.split("\n")[0]
def get_conda_packages(run_lambda, patterns=None):
if patterns is None:
patterns = DEFAULT_CONDA_PATTERNS
conda = os.environ.get("CONDA_EXE", "conda")
out = run_and_read_all(run_lambda, [conda, "list"])
if out is None:
return out
return "\n".join(
line for line in out.splitlines() if not line.startswith("#") and any(name in line for name in patterns)
)
def get_gcc_version(run_lambda):
return run_and_parse_first_match(run_lambda, "gcc --version", r"gcc (.*)")
def get_clang_version(run_lambda):
return run_and_parse_first_match(run_lambda, "clang --version", r"clang version (.*)")
def get_cmake_version(run_lambda):
return run_and_parse_first_match(run_lambda, "cmake --version", r"cmake (.*)")
def get_nvidia_driver_version(run_lambda):
if get_platform() == "darwin":
cmd = "kextstat | grep -i cuda"
return run_and_parse_first_match(run_lambda, cmd, r"com[.]nvidia[.]CUDA [(](.*?)[)]")
smi = get_nvidia_smi()
return run_and_parse_first_match(run_lambda, smi, r"Driver Version: (.*?) ")
def get_gpu_info(run_lambda):
if get_platform() == "darwin" or (
TORCH_AVAILABLE and hasattr(torch.version, "hip") and torch.version.hip is not None
):
if TORCH_AVAILABLE and torch.cuda.is_available():
if torch.version.hip is not None:
prop = torch.cuda.get_device_properties(0)
if hasattr(prop, "gcnArchName"):
gcnArch = " ({})".format(prop.gcnArchName)
else:
gcnArch = "NoGCNArchNameOnOldPyTorch"
else:
gcnArch = ""
return torch.cuda.get_device_name(None) + gcnArch
return None
smi = get_nvidia_smi()
uuid_regex = re.compile(r" \(UUID: .+?\)")
rc, out, _ = run_lambda(smi + " -L")
if rc != 0:
return None
# Anonymize GPUs by removing their UUID
return re.sub(uuid_regex, "", out)
def get_running_cuda_version(run_lambda):
return run_and_parse_first_match(run_lambda, "nvcc --version", r"release .+ V(.*)")
def get_cudnn_version(run_lambda):
"""Return a list of libcudnn.so; it's hard to tell which one is being used."""
if get_platform() == "win32":
system_root = os.environ.get("SYSTEMROOT", "C:\\Windows")
cuda_path = os.environ.get("CUDA_PATH", "%CUDA_PATH%")
where_cmd = os.path.join(system_root, "System32", "where")
cudnn_cmd = '{} /R "{}\\bin" cudnn*.dll'.format(where_cmd, cuda_path)
elif get_platform() == "darwin":
# CUDA libraries and drivers can be found in /usr/local/cuda/. See
# https://docs.nvidia.com/cuda/cuda-installation-guide-mac-os-x/index.html#install
# https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#installmac
# Use CUDNN_LIBRARY when cudnn library is installed elsewhere.
cudnn_cmd = "ls /usr/local/cuda/lib/libcudnn*"
else:
cudnn_cmd = 'ldconfig -p | grep libcudnn | rev | cut -d" " -f1 | rev'
rc, out, _ = run_lambda(cudnn_cmd)
# find will return 1 if there are permission errors or if not found
if len(out) == 0 or (rc != 1 and rc != 0):
l = os.environ.get("CUDNN_LIBRARY")
if l is not None and os.path.isfile(l):
return os.path.realpath(l)
return None
files_set = set()
for fn in out.split("\n"):
fn = os.path.realpath(fn) # eliminate symbolic links
if os.path.isfile(fn):
files_set.add(fn)
if not files_set:
return None
# Alphabetize the result because the order is non-deterministic otherwise
files = sorted(files_set)
if len(files) == 1:
return files[0]
result = "\n".join(files)
return "Probably one of the following:\n{}".format(result)
def get_nvidia_smi():
# Note: nvidia-smi is currently available only on Windows and Linux
smi = "nvidia-smi"
if get_platform() == "win32":
system_root = os.environ.get("SYSTEMROOT", "C:\\Windows")
program_files_root = os.environ.get("PROGRAMFILES", "C:\\Program Files")
legacy_path = os.path.join(program_files_root, "NVIDIA Corporation", "NVSMI", smi)
new_path = os.path.join(system_root, "System32", smi)
smis = [new_path, legacy_path]
for candidate_smi in smis:
if os.path.exists(candidate_smi):
smi = '"{}"'.format(candidate_smi)
break
return smi
def get_rocm_version(run_lambda):
"""Returns the ROCm version if available, otherwise 'N/A'."""
return run_and_parse_first_match(run_lambda, "hipcc --version", r"HIP version: (\S+)")
def get_vllm_version():
from vllm import __version__, __version_tuple__
if __version__ == "dev":
return "N/A (dev)"
version_str = __version_tuple__[-1]
if isinstance(version_str, str) and version_str.startswith("g"):
# it's a dev build
if "." in version_str:
# it's a dev build containing local changes
git_sha = version_str.split(".")[0][1:]
date = version_str.split(".")[-1][1:]
return f"{__version__} (git sha: {git_sha}, date: {date})"
else:
# it's a dev build without local changes
git_sha = version_str[1:] # type: ignore
return f"{__version__} (git sha: {git_sha})"
return __version__
def get_vllm_omni_version(run_lambda):
try:
import vllm_omni
from vllm_omni import __version__, __version_tuple__
version_str = __version_tuple__[-1]
if isinstance(version_str, str) and version_str.startswith("g"):
if "." in version_str:
git_sha = version_str.split(".")[0][1:]
date = version_str.split(".")[-1][1:]
return f"{__version__} (git sha: {git_sha}, date: {date})"
else:
git_sha = version_str[1:]
return f"{__version__} (git sha: {git_sha})"
package_dir = os.path.dirname(os.path.abspath(vllm_omni.__file__))
git_sha = run_and_read_all(run_lambda, f"git -C {package_dir} rev-parse --short HEAD")
if git_sha:
return f"{__version__} (git sha: {git_sha})"
return __version__
except ImportError:
return "N/A (vllm_omni not installed)"
def summarize_vllm_build_flags():
# This could be a static method if the flags are constant, or dynamic if you need to check environment variables, etc.
return "CUDA Archs: {}; ROCm: {}".format(
os.environ.get("TORCH_CUDA_ARCH_LIST", "Not Set"),
"Enabled" if os.environ.get("ROCM_HOME") else "Disabled",
)
def get_gpu_topo(run_lambda):
output = None
if get_platform() == "linux":
output = run_and_read_all(run_lambda, "nvidia-smi topo -m")
if output is None:
output = run_and_read_all(run_lambda, "rocm-smi --showtopo")
return output
def get_cpu_info(run_lambda):
rc, out, err = 0, "", ""
if get_platform() == "linux":
rc, out, err = run_lambda("lscpu")
elif get_platform() == "win32":
rc, out, err = run_lambda(
"wmic cpu get Name,Manufacturer,Family,Architecture,ProcessorType,DeviceID, \
CurrentClockSpeed,MaxClockSpeed,L2CacheSize,L2CacheSpeed,Revision /VALUE"
)
elif get_platform() == "darwin":
rc, out, err = run_lambda("sysctl -n machdep.cpu.brand_string")
cpu_info = "None"
if rc == 0:
cpu_info = out
else:
cpu_info = err
return cpu_info
def get_platform():
if sys.platform.startswith("linux"):
return "linux"
elif sys.platform.startswith("win32"):
return "win32"
elif sys.platform.startswith("cygwin"):
return "cygwin"
elif sys.platform.startswith("darwin"):
return "darwin"
else:
return sys.platform
def get_mac_version(run_lambda):
return run_and_parse_first_match(run_lambda, "sw_vers -productVersion", r"(.*)")
def get_windows_version(run_lambda):
system_root = os.environ.get("SYSTEMROOT", "C:\\Windows")
wmic_cmd = os.path.join(system_root, "System32", "Wbem", "wmic")
findstr_cmd = os.path.join(system_root, "System32", "findstr")
return run_and_read_all(run_lambda, "{} os get Caption | {} /v Caption".format(wmic_cmd, findstr_cmd))
def get_lsb_version(run_lambda):
return run_and_parse_first_match(run_lambda, "lsb_release -a", r"Description:\t(.*)")
def check_release_file(run_lambda):
return run_and_parse_first_match(run_lambda, "cat /etc/*-release", r'PRETTY_NAME="(.*)"')
def get_os(run_lambda):
from platform import machine
platform = get_platform()
if platform == "win32" or platform == "cygwin":
return get_windows_version(run_lambda)
if platform == "darwin":
version = get_mac_version(run_lambda)
if version is None:
return None
return "macOS {} ({})".format(version, machine())
if platform == "linux":
# Ubuntu/Debian based
desc = get_lsb_version(run_lambda)
if desc is not None:
return "{} ({})".format(desc, machine())
# Try reading /etc/*-release
desc = check_release_file(run_lambda)
if desc is not None:
return "{} ({})".format(desc, machine())
return "{} ({})".format(platform, machine())
# Unknown platform
return platform
def get_python_platform():
import platform
return platform.platform()
def get_libc_version():
import platform
if get_platform() != "linux":
return "N/A"
return "-".join(platform.libc_ver())
def is_uv_venv():
if os.environ.get("UV"):
return True
pyvenv_cfg_path = os.path.join(sys.prefix, "pyvenv.cfg")
if os.path.exists(pyvenv_cfg_path):
with open(pyvenv_cfg_path, "r") as f:
return any(line.startswith("uv = ") for line in f)
return False
def get_pip_packages(run_lambda, patterns=None):
"""Return `pip list` output. Note: will also find conda-installed pytorch and numpy packages."""
if patterns is None:
patterns = DEFAULT_PIP_PATTERNS
def run_with_pip():
try:
import importlib.util
pip_spec = importlib.util.find_spec("pip")
pip_available = pip_spec is not None
except ImportError:
pip_available = False
if pip_available:
cmd = [sys.executable, "-mpip", "list", "--format=freeze"]
elif is_uv_venv():
print("uv is set")
cmd = ["uv", "pip", "list", "--format=freeze"]
else:
raise RuntimeError("Could not collect pip list output (pip or uv module not available)")
out = run_and_read_all(run_lambda, cmd)
return "\n".join(line for line in out.splitlines() if any(name in line for name in patterns))
pip_version = "pip3" if sys.version[0] == "3" else "pip"
out = run_with_pip()
return pip_version, out
def get_cachingallocator_config():
ca_config = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
return ca_config
def get_cuda_module_loading_config():
if TORCH_AVAILABLE and torch.cuda.is_available():
torch.cuda.init()
config = os.environ.get("CUDA_MODULE_LOADING", "")
return config
else:
return "N/A"
def is_xnnpack_available():
if TORCH_AVAILABLE:
import torch.backends.xnnpack
return str(torch.backends.xnnpack.enabled) # type: ignore[attr-defined]
else:
return "N/A"
def get_env_vars():
env_vars = ""
secret_terms = ("secret", "token", "api", "access", "password")
report_prefix = (
"TORCH",
"NCCL",
"PYTORCH",
"CUDA",
"CUBLAS",
"CUDNN",
"OMP_",
"MKL_",
"NVIDIA",
)
for k, v in os.environ.items():
if any(term in k.lower() for term in secret_terms):
continue
if k in environment_variables:
env_vars = env_vars + "{}={}".format(k, v) + "\n"
if k.startswith(report_prefix):
env_vars = env_vars + "{}={}".format(k, v) + "\n"
return env_vars
def get_env_info():
run_lambda = run
pip_version, pip_list_output = get_pip_packages(run_lambda)
if TORCH_AVAILABLE:
version_str = torch.__version__
debug_mode_str = str(torch.version.debug)
cuda_available_str = str(torch.cuda.is_available())
cuda_version_str = torch.version.cuda
if not hasattr(torch.version, "hip") or torch.version.hip is None: # cuda version
hip_compiled_version = hip_runtime_version = miopen_runtime_version = "N/A"
else: # HIP version
def get_version_or_na(cfg, prefix):
_lst = [s.rsplit(None, 1)[-1] for s in cfg if prefix in s]
return _lst[0] if _lst else "N/A"
cfg = torch._C._show_config().split("\n")
hip_runtime_version = get_version_or_na(cfg, "HIP Runtime")
miopen_runtime_version = get_version_or_na(cfg, "MIOpen")
cuda_version_str = "N/A"
hip_compiled_version = torch.version.hip
else:
version_str = debug_mode_str = cuda_available_str = cuda_version_str = "N/A"
hip_compiled_version = hip_runtime_version = miopen_runtime_version = "N/A"
sys_version = sys.version.replace("\n", " ")
conda_packages = get_conda_packages(run_lambda)
rocm_version = get_rocm_version(run_lambda)
vllm_version = get_vllm_version()
vllm_omni_version = get_vllm_omni_version(run_lambda)
vllm_build_flags = summarize_vllm_build_flags()
gpu_topo = get_gpu_topo(run_lambda)
return SystemEnv(
torch_version=version_str,
is_debug_build=debug_mode_str,
python_version="{} ({}-bit runtime)".format(sys_version, sys.maxsize.bit_length() + 1),
python_platform=get_python_platform(),
is_cuda_available=cuda_available_str,
cuda_compiled_version=cuda_version_str,
cuda_runtime_version=get_running_cuda_version(run_lambda),
cuda_module_loading=get_cuda_module_loading_config(),
nvidia_gpu_models=get_gpu_info(run_lambda),
nvidia_driver_version=get_nvidia_driver_version(run_lambda),
cudnn_version=get_cudnn_version(run_lambda),
hip_compiled_version=hip_compiled_version,
hip_runtime_version=hip_runtime_version,
miopen_runtime_version=miopen_runtime_version,
pip_version=pip_version,
pip_packages=pip_list_output,
conda_packages=conda_packages,
os=get_os(run_lambda),
libc_version=get_libc_version(),
gcc_version=get_gcc_version(run_lambda),
clang_version=get_clang_version(run_lambda),
cmake_version=get_cmake_version(run_lambda),
caching_allocator_config=get_cachingallocator_config(),
is_xnnpack_available=is_xnnpack_available(),
cpu_info=get_cpu_info(run_lambda),
rocm_version=rocm_version,
vllm_version=vllm_version,
vllm_omni_version=vllm_omni_version,
vllm_build_flags=vllm_build_flags,
gpu_topo=gpu_topo,
env_vars=get_env_vars(),
)
env_info_fmt = """
==============================
System Info
==============================
OS : {os}
GCC version : {gcc_version}
Clang version : {clang_version}
CMake version : {cmake_version}
Libc version : {libc_version}
==============================
PyTorch Info
==============================
PyTorch version : {torch_version}
Is debug build : {is_debug_build}
CUDA used to build PyTorch : {cuda_compiled_version}
ROCM used to build PyTorch : {hip_compiled_version}
==============================
Python Environment
==============================
Python version : {python_version}
Python platform : {python_platform}
==============================
CUDA / GPU Info
==============================
Is CUDA available : {is_cuda_available}
CUDA runtime version : {cuda_runtime_version}
CUDA_MODULE_LOADING set to : {cuda_module_loading}
GPU models and configuration : {nvidia_gpu_models}
Nvidia driver version : {nvidia_driver_version}
cuDNN version : {cudnn_version}
HIP runtime version : {hip_runtime_version}
MIOpen runtime version : {miopen_runtime_version}
Is XNNPACK available : {is_xnnpack_available}
==============================
CPU Info
==============================
{cpu_info}
==============================
Versions of relevant libraries
==============================
{pip_packages}
{conda_packages}
""".strip()
# both the above code and the following code use `strip()` to
# remove leading/trailing whitespaces, so we need to add a newline
# in between to separate the two sections
env_info_fmt += "\n\n"
env_info_fmt += """
==============================
vLLM Info
==============================
ROCM Version : {rocm_version}
vLLM Version : {vllm_version}
vLLM-Omni Version : {vllm_omni_version}
vLLM Build Flags:
{vllm_build_flags}
GPU Topology:
{gpu_topo}
==============================
Environment Variables
==============================
{env_vars}
""".strip()
def pretty_str(envinfo):
def replace_nones(dct, replacement="Could not collect"):
for key in dct.keys():
if dct[key] is not None:
continue
dct[key] = replacement
return dct
def replace_bools(dct, true="Yes", false="No"):
for key in dct.keys():
if dct[key] is True:
dct[key] = true
elif dct[key] is False:
dct[key] = false
return dct
def prepend(text, tag="[prepend]"):
lines = text.split("\n")
updated_lines = [tag + line for line in lines]
return "\n".join(updated_lines)
def replace_if_empty(text, replacement="No relevant packages"):
if text is not None and len(text) == 0:
return replacement
return text
def maybe_start_on_next_line(string):
# If `string` is multiline, prepend a \n to it.
if string is not None and len(string.split("\n")) > 1:
return "\n{}\n".format(string)
return string
mutable_dict = envinfo._asdict()
# If nvidia_gpu_models is multiline, start on the next line
mutable_dict["nvidia_gpu_models"] = maybe_start_on_next_line(envinfo.nvidia_gpu_models)
# If the machine doesn't have CUDA, report some fields as 'No CUDA'
dynamic_cuda_fields = [
"cuda_runtime_version",
"nvidia_gpu_models",
"nvidia_driver_version",
]
all_cuda_fields = dynamic_cuda_fields + ["cudnn_version"]
all_dynamic_cuda_fields_missing = all(mutable_dict[field] is None for field in dynamic_cuda_fields)
if TORCH_AVAILABLE and not torch.cuda.is_available() and all_dynamic_cuda_fields_missing:
for field in all_cuda_fields:
mutable_dict[field] = "No CUDA"
if envinfo.cuda_compiled_version is None:
mutable_dict["cuda_compiled_version"] = "None"
# Replace True with Yes, False with No
mutable_dict = replace_bools(mutable_dict)
# Replace all None objects with 'Could not collect'
mutable_dict = replace_nones(mutable_dict)
# If either of these are '', replace with 'No relevant packages'
mutable_dict["pip_packages"] = replace_if_empty(mutable_dict["pip_packages"])
mutable_dict["conda_packages"] = replace_if_empty(mutable_dict["conda_packages"])
# Tag conda and pip packages with a prefix
# If they were previously None, they'll show up as ie '[conda] Could not collect'
if mutable_dict["pip_packages"]:
mutable_dict["pip_packages"] = prepend(mutable_dict["pip_packages"], "[{}] ".format(envinfo.pip_version))
if mutable_dict["conda_packages"]:
mutable_dict["conda_packages"] = prepend(mutable_dict["conda_packages"], "[conda] ")
mutable_dict["cpu_info"] = envinfo.cpu_info
return env_info_fmt.format(**mutable_dict)
def get_pretty_env_info():
return pretty_str(get_env_info())
def main():
print("Collecting environment information...")
output = get_pretty_env_info()
print(output)
if TORCH_AVAILABLE and hasattr(torch, "utils") and hasattr(torch.utils, "_crash_handler"):
minidump_dir = torch.utils._crash_handler.DEFAULT_MINIDUMP_DIR
if sys.platform == "linux" and os.path.exists(minidump_dir):
dumps = [os.path.join(minidump_dir, dump) for dump in os.listdir(minidump_dir)]
latest = max(dumps, key=os.path.getctime)
ctime = os.path.getctime(latest)
creation_time = datetime.datetime.fromtimestamp(ctime).strftime("%Y-%m-%d %H:%M:%S")
msg = (
"\n*** Detected a minidump at {} created on {}, ".format(latest, creation_time)
+ "if this is related to your bug please include it when you file a report ***"
)
print(msg, file=sys.stderr)
if __name__ == "__main__":
main()
+20
View File
@@ -0,0 +1,20 @@
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
ARG VLLM_BASE_TAG=v0.14.0
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
ARG APP_DIR=/workspace/vllm-omni
WORKDIR ${APP_DIR}
COPY . .
# Install system dependencies
RUN apt-get update && \
apt-get install -y ffmpeg && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
# Install vllm-omni into the same uv-managed Python environment used by the base image.
RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
RUN ln -sf /usr/bin/python3 /usr/bin/python
ENTRYPOINT []
+15
View File
@@ -0,0 +1,15 @@
ARG VLLM_ASCEND_IMAGE=quay.nju.edu.cn/ascend/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.11.0rc2
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
ARG APP_DIR=/vllm-workspace/vllm-omni
WORKDIR ${APP_DIR}
COPY . .
# Install vllm-omni with dev dependencies
RUN pip install --no-cache-dir -e ".[dev]"
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
ENTRYPOINT []
+32
View File
@@ -0,0 +1,32 @@
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.14.0
FROM ${BASE_IMAGE} AS final
ARG COMMON_WORKDIR=/app
WORKDIR ${COMMON_WORKDIR}
# Step 1: Setup - Install system dependencies
RUN apt-get update && \
apt-get install -y ffmpeg && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
# Step 2: Copy vllm-omni code and install without uv
COPY . ${COMMON_WORKDIR}/vllm-omni
RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
# When we are installing onnxruntime-rocm, we need to uninstall the system-installed onnxruntime first.
# These are the dependencies of Qwen3-TTS.
RUN uv pip uninstall onnxruntime --system && uv pip install --no-cache-dir onnxruntime-rocm sox --system
RUN ln -sf /usr/bin/python3 /usr/bin/python
CMD ["/bin/bash"]
ENTRYPOINT []
#Set entrypoint for vllm-openai official images
FROM final AS vllm-openai
ENTRYPOINT ["vllm", "serve", "--omni"]
+69
View File
@@ -0,0 +1,69 @@
nav:
- Home: README.md
- User Guide:
- Getting Started:
- getting_started/quickstart.md
- getting_started/installation/*
- Serving:
- OpenAI-Compatible API:
- Image Generation: serving/image_generation_api.md
- Examples:
- examples/README.md
- Offline Inference:
- Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
- Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
- Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
- Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
- Qwen3-TTS Offline Inference: user_guide/examples/offline_inference/qwen3_tts.md
- Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
- Online Serving:
- Image-To-Image: user_guide/examples/online_serving/image_to_image.md
- Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
- Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
- Text-To-Image: user_guide/examples/online_serving/text_to_image.md
- General:
- usage/*
- Configuration:
- configuration/README.md
- configuration/*
- Models:
- models/supported_models.md
- Features:
- Sleep Mode: features/sleep_mode.md
- Diffusion Features:
- Overview: user_guide/diffusion_acceleration.md
- TeaCache: user_guide/diffusion/teacache.md
- Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
- Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
- CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
- Developer Guide:
- General:
- contributing/README.md
- glob: contributing/*
flatten_single_child_sections: true
- Model Implementation:
- contributing/model/README.md
- contributing/model/adding_omni_model.md
- contributing/model/adding_diffusion_model.md
- CI: contributing/ci
- Design Documents:
- design/index.md
- design/architecture_overview.md
- Feature Design:
- design/feature/disaggregated_inference.md
- design/feature/ray_based_execution.md
- Module Design:
- design/module/ar_module.md
- design/module/dit_module.md
- design/module/entrypoint_module.md
- Docs Guide: contributing/DOCS_GUIDE.md
- API Reference:
- api/README.md
- api/vllm_omni
- CLI Reference: cli
- Community:
- community/*
- Slack: https://slack.vllm.ai
- Blog: https://blog.vllm.ai
- Forum: https://discuss.vllm.ai
+64
View File
@@ -0,0 +1,64 @@
---
hide:
- navigation
- toc
---
# Welcome to vLLM-Omni
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="./source/logos/vllm-omni-logo.png">
<img alt="vllm-omni" src="./source/logos/vllm-omni-logo.png" width=55%>
</picture>
</p>
<h3 align="center">
Easy, fast, and cheap omni-modality model serving for everyone
</h3>
<p style="text-align:center">
<script async defer src="https://buttons.github.io/buttons.js"></script>
<a class="github-button" href="https://github.com/vllm-project/vllm-omni" data-show-count="true" data-size="large" aria-label="Star">Star</a>
<a class="github-button" href="https://github.com/vllm-project/vllm-omni/subscription" data-show-count="true" data-icon="octicon-eye" data-size="large" aria-label="Watch">Watch</a>
<a class="github-button" href="https://github.com/vllm-project/vllm-omni/fork" data-show-count="true" data-icon="octicon-repo-forked" data-size="large" aria-label="Fork">Fork</a>
</p>
## About
[vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni is a framework that extends its support for omni-modality model inference and serving:
- **Omni-modality**: Text, image, video, and audio data processing
- **Non-autoregressive Architectures**: extend the AR support of vLLM to Diffusion Transformers (DiT) and other parallel generation models
- **Heterogeneous outputs**: from traditional text generation to multimodal outputs
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="./source/architecture/omni-modality-model-architecture.png">
<img alt="vllm-omni-arch" src="./source/architecture/omni-modality-model-architecture.png" width=55%>
</picture>
</p>
vLLM-Omni is fast with:
- State-of-the-art AR support by leveraging efficient KV cache management from vLLM
- Pipelined stage execution overlapping for high throughput performance
- Fully disaggregation based on OmniConnector and dynamic resource allocation across stages
vLLM-Omni is flexible and easy to use with:
- Heterogeneous pipeline abstraction to manage complex model workflows
- Seamless integration with popular Hugging Face models
- Tensor, pipeline, data and expert parallelism support for distributed inference
- Streaming outputs
- OpenAI-compatible API server
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
- Omni-modality models (e.g. Qwen2.5-Omni, Qwen3-Omni)
- Multi-modality generation models (e.g. Qwen-Image)
For more information, checkout the following:
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true)
- [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
+107
View File
@@ -0,0 +1,107 @@
# Summary
## Entry Points
Main entry points for vLLM-Omni inference and serving.
- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
- [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalContentParser][]
- [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalItemTracker][]
- [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][]
- [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
- [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
- [vllm_omni.entrypoints.log_utils.OrchestratorMetrics][]
- [vllm_omni.entrypoints.log_utils.StageRequestMetrics][]
- [vllm_omni.entrypoints.log_utils.StageStats][]
- [vllm_omni.entrypoints.omni.Omni][]
- [vllm_omni.entrypoints.omni.OmniBase][]
- [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][]
- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
- [vllm_omni.entrypoints.omni_stage.OmniStage][]
- [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
## Inputs
Input data structures for multi-modal inputs.
- [vllm_omni.inputs.data.OmniEmbedsPrompt][]
- [vllm_omni.inputs.data.OmniTokenInputs][]
- [vllm_omni.inputs.data.OmniTokensPrompt][]
- [vllm_omni.inputs.parse.parse_singleton_prompt_omni][]
- [vllm_omni.inputs.preprocess.OmniInputPreprocessor][]
## Engine
Engine classes for offline and online inference.
- [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
- [vllm_omni.engine.AdditionalInformationEntry][]
- [vllm_omni.engine.AdditionalInformationPayload][]
- [vllm_omni.engine.OmniEngineCoreOutput][]
- [vllm_omni.engine.OmniEngineCoreOutputs][]
- [vllm_omni.engine.OmniEngineCoreRequest][]
- [vllm_omni.engine.PromptEmbedsPayload][]
- [vllm_omni.engine.arg_utils.AsyncOmniEngineArgs][]
- [vllm_omni.engine.arg_utils.OmniEngineArgs][]
- [vllm_omni.engine.input_processor.OmniInputProcessor][]
- [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
- [vllm_omni.engine.output_processor.OmniRequestState][]
## Core
Core scheduling and caching components.
- [vllm_omni.core.sched.omni_ar_scheduler.KVCacheTransferData][]
- [vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler][]
- [vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler][]
- [vllm_omni.core.sched.output.OmniCachedRequestData][]
- [vllm_omni.core.sched.output.OmniNewRequestData][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedGroupResidualVectorQuantization][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedResidualVectorQuantization][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.EuclideanCodebook][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.VectorQuantization][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.preprocess][]
## Configuration
Configuration classes.
- [vllm_omni.config.model.OmniModelConfig][]
- [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
- [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
- [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSSpeakerEncoderConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSTalkerCodePredictorConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSTalkerConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2.Qwen3TTSTokenizerV2Config][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2.Qwen3TTSTokenizerV2DecoderConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1Config][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderBigVGANConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
## Workers
Worker classes and model runners for distributed inference.
- [vllm_omni.diffusion.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]
- [vllm_omni.diffusion.worker.gpu_diffusion_worker.GPUDiffusionWorker][]
- [vllm_omni.diffusion.worker.gpu_diffusion_worker.WorkerProc][]
- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorker][]
- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorkerProc][]
- [vllm_omni.worker.gpu_ar_model_runner.ExecuteModelState][]
- [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
- [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
- [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][]
- [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
- [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
- [vllm_omni.worker.npu.npu_ar_model_runner.ExecuteModelState][]
- [vllm_omni.worker.npu.npu_ar_model_runner.NPUARModelRunner][]
- [vllm_omni.worker.npu.npu_ar_worker.NPUARWorker][]
- [vllm_omni.worker.npu.npu_generation_model_runner.NPUGenerationModelRunner][]
- [vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker][]
- [vllm_omni.worker.npu.npu_model_runner.OmniNPUModelRunner][]
Binary file not shown.

After

Width:  |  Height:  |  Size: 126 KiB

+24
View File
@@ -0,0 +1,24 @@
# vLLM-Omni CLI Guide
The CLI for vLLM-Omni inherits from vllm with some additional arguments.
## serve
Starts the vLLM-Omni OpenAI Compatible API server.
Start with a model:
```bash
vllm serve Qwen/Qwen2.5-Omni-7B --omni
```
Specify the port:
```bash
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
```
If you have custom stage configs file, launch the server with command below
```bash
vllm serve Qwen/Qwen2.5-Omni-7B --omni --stage-configs-path /path/to/stage_configs_file
```
+5
View File
@@ -0,0 +1,5 @@
# 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
+1
View File
@@ -0,0 +1 @@
# Meetups
+12
View File
@@ -0,0 +1,12 @@
# 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!
+21
View File
@@ -0,0 +1,21 @@
# Configuration Options
This section lists the most common options for running vLLM-Omni.
For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html)
Currently, the main options are maintained by stage configs for each model.
For specific example, please refer to [Qwen2.5-omni stage config](stage_configs/qwen2_5_omni.yaml)
For introduction, please check [Introduction for stage config](./stage_configs.md)
## Memory Configuration
- **[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
## Optimization Features
- **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
- **[Cache-DiT Configuration](../user_guide/diffusion/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models
- **[Parallelism Configuration](../user_guide/diffusion/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models
@@ -0,0 +1,207 @@
# GPU Memory Calculation and Configuration
This guide explains how to calculate GPU memory requirements and properly configure `gpu_memory_utilization` for vLLM-Omni stages.
## Overview
`gpu_memory_utilization` is a critical parameter that controls how much GPU memory each stage can use. It's specified as a fraction between 0.0 and 1.0, where:
- `0.8` means 80% of the GPU's total memory
- `1.0` means 100% of the GPU's total memory (not recommended, leaves no buffer)
## How Memory is Calculated
### Memory Allocation Formula
For each stage, vLLM-Omni calculates the requested memory as:
```
requested_memory = total_gpu_memory × gpu_memory_utilization
```
The system checks that:
```
free_memory ≥ requested_memory
```
If this condition is not met, the stage will fail to initialize with an error message showing the memory requirements.
### Memory Components
The total memory used by a stage includes:
1. **Model Weights**: The size of the model parameters loaded on the GPU
2. **KV Cache**: Memory for storing key-value cache during generation
3. **Activation Memory**: Temporary memory for intermediate computations
4. **System Overhead**: Memory used by CUDA, PyTorch, and other system components
5. **Non-Torch Memory**: Memory allocated outside of PyTorch (e.g., CUDA graphs)
### Example Calculation
For a GPU with 80GB total memory:
- `gpu_memory_utilization: 0.8` → 64GB available for the stage
- `gpu_memory_utilization: 0.6` → 48GB available for the stage
- `gpu_memory_utilization: 0.15` → 12GB available for the stage
## Setting Up `gpu_memory_utilization`
### Step 1: Determine GPU Memory
First, check your GPU's total memory:
```bash
# Using nvidia-smi
nvidia-smi --query-gpu=memory.total --format=csv
# Or using Python
python -c "import torch; print(f'{torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')"
```
### Step 2: Estimate Model Memory Requirements
#### For Autoregressive (AR) Stages
AR stages typically need more memory due to:
- Large model weights
- KV cache for attention
- Activation buffers
#### For Diffusion/Generation Stages
Diffusion stages (like code2wav) typically need less memory:
- Smaller model components
- Different memory access patterns
**Typical values:**
- `0.1 - 0.3` for most diffusion stages
### Step 3: Consider Multi-Stage Scenarios
When multiple stages share the same GPU, you must ensure the sum of their `gpu_memory_utilization` values doesn't exceed 1.0.
**Example: Two stages on GPU 0**
```yaml
stage_args:
- stage_id: 0
runtime:
devices: "0"
engine_args:
gpu_memory_utilization: 0.6 # Uses 60% of GPU 0
- stage_id: 1
runtime:
devices: "0"
engine_args:
gpu_memory_utilization: 0.3 # Uses 30% of GPU 0
# Total: 90% of GPU 0 (safe, leaves 10% buffer)
```
**Important:** If stages run on different GPUs, each can use up to 1.0 independently.
### Step 4: Account for Tensor Parallelism
When using `tensor_parallel_size > 1`, the model is split across multiple GPUs, so each GPU needs less memory.
**Example: 2-way tensor parallelism**
```yaml
stage_args:
- stage_id: 0
runtime:
devices: "0,1" # Uses both GPUs
engine_args:
tensor_parallel_size: 2
gpu_memory_utilization: 0.6 # 60% per GPU
# Model is split, so each GPU uses ~30% of model memory
```
## Examples
### Qwen3-Omni-MoE on 2x H100-80GB
```yaml
stage_args:
- stage_id: 0 # Thinker stage with TP=2
runtime:
devices: "0,1"
engine_args:
tensor_parallel_size: 2
gpu_memory_utilization: 0.6 # 48GB per GPU
- stage_id: 1 # Talker stage
runtime:
devices: "1"
engine_args:
gpu_memory_utilization: 0.3 # 24GB on GPU 1
- stage_id: 2 # Code2Wav stage
runtime:
devices: "0"
engine_args:
gpu_memory_utilization: 0.1 # 8GB on GPU 0
```
**Note:** In this configuration, stages 0 and 2 share GPU 0, but they run at different times in the pipeline, so their memory usage doesn't overlap.
## Troubleshooting
### Error: "Free memory is less than desired GPU memory utilization"
This means the GPU doesn't have enough free memory when the stage starts.
**Solutions:**
1. Free up memory by closing other processes
2. Reduce `gpu_memory_utilization` for this stage
3. Use a GPU with more memory
4. Move the stage to a different GPU
### Error: OOM during inference
The stage initialized but ran out of memory during processing.
**Solutions:**
1. Reduce `max_num_batched_tokens`
2. Reduce `max_batch_size` in runtime config
3. Lower `gpu_memory_utilization` slightly
4. Enable quantization if supported
### Memory Not Fully Utilized
If you see low memory usage, you can:
1. Increase `gpu_memory_utilization` to allow larger KV cache
2. Increase `max_num_batched_tokens` for better batching
3. Check if other stages are limiting throughput
## Useful formula for Memory Calculation
### KV Cache Memory
The KV cache size depends on:
- Number of sequences in batch
- Sequence length (prompt + generation)
- Model hidden size
- Number of attention heads
- Number of layers
approximate Formula:
```
kv_cache_memory ≈ batch_size × seq_len × hidden_size × num_layers × 2 × dtype_size
```
2 for k & v
### Model Weight Memory
```
model_memory ≈ num_parameters × dtype_size
```
For example:
- 7B parameters in FP16: ~14GB
- 7B parameters in FP32: ~28GB
- 7B parameters in INT8: ~7GB
### Activation Memory
Activation memory is typically smaller but varies with:
- Batch size
- Sequence length
- Model architecture
It's usually 10-30% of model weight memory during inference.
+275
View File
@@ -0,0 +1,275 @@
# Stage configs for vLLM-Omni
In vLLM-Omni, the target model is separated into multiple stages, which are processed by different LLMEngines, DiffusionEngines or other types of engines. Depending on different types of stages, such as Autoregressive (AR) stage or Diffusion transformer (DiT) stage, each can choose corresponding schedulers, model workers to load with the Engines in a plug-in fashion.
!!! note
Default stage config YAMLs (for example, `vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml` and `vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml`) are bundled and loaded automatically when `stage_configs_path` is not provided. They have been verified to work on 1xH100 for Qwen2.5-Omni and 2xH100 for Qwen3-Omni.
Therefore, as a core part of vLLM-Omni, the stage configs for a model have several main functions:
- Claim partition of stages and their corresponding class implementation in `model_executor/models`.
- The disaggregated configuration for each stage and the communication topology among them.
- Engine arguments for each engine within the stage.
- Input and output dependencies for each stage.
- Default input parameters.
If users want to modify some part of it. The custom stage_configs file can be input as input argument in both online and offline. Just like examples below:
For offline (Assume necessary dependencies have ben imported):
```python
model_name = "Qwen/Qwen2.5-Omni-7B"
omni_llm = OmniLLM(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
```
For online serving:
```bash
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
```
!!! important
We are actively iterating on the definition of stage configs, and we welcome all feedbacks from both community users and developers to help us shape the development!
Below is a specific example of stage_configs.yaml in Qwen2.5-omni.
```python
# stage config for running qwen2.5-omni with architecture of OmniLLM.
stage_args:
- stage_id: 0 # mark the unique id for each stage
runtime: # The disaggregated configuration
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1 # the batch_size for offline inference
engine_args: # Engine arguments for a certain engine
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker # The specific worker used
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler # The specific scehduler used
gpu_memory_utilization: 0.8 # The gpu memory allocation for the stage within a single chip
enforce_eager: true # Now we only support eager mode
trust_remote_code: true # Needed by huggingface config parsing
engine_output_type: latent # It claims that the stage will input latent hiddenstates besides token ids
enable_prefix_caching: false # For request with hiddenstates output, the prefix caching is not supported now
is_comprehension: true # If the stage is a text or multimodal comprehension module. If it is, the AsyncOmni will use its tokenizer as default
final_output: true # If the stage has output as part of final outputs. If it is false, which means that the stage only works as a intermediate role.
final_output_type: text # What is the final output type. It can be text and audio now.
default_sampling_params: # sampling parameters for the stage. Their meaning aligns with vLLM.
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 2048
seed: 42
detokenize: True
repetition_penalty: 1.1
- stage_id: 1
runtime:
process: true
devices: "1"
max_batch_size: 3
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.8
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
engine_output_type: latent
engine_input_source: [0]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
default_sampling_params:
temperature: 0.9
top_p: 0.8
top_k: 40
max_tokens: 2048
seed: 42
detokenize: True
repetition_penalty: 1.05
stop_token_ids: [8294]
- stage_id: 2
runtime:
process: true
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
max_batch_size: 1
engine_args:
model_stage: code2wav
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
gpu_memory_utilization: 0.15
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
engine_output_type: audio
engine_input_source: [1]
final_output: true
final_output_type: audio
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 2048
seed: 42
detokenize: True
repetition_penalty: 1.1
# Top-level runtime config (concise): default windows and stage edges
runtime:
enabled: true
defaults:
window_size: -1 # Simplified: trigger downstream only after full upstream completion
max_inflight: 1 # Simplified: process serially within each stage
edges:
- from: 0 # thinker → talker: trigger only after receiving full input (-1)
to: 1
window_size: -1
- from: 1 # talker → code2wav: trigger only after receiving full input (-1)
to: 2
window_size: -1
```
## Stage Configuration Arguments
Each stage in the `stage_args` list contains the following configuration options:
### `stage_id`
A unique identifier for each stage in the multi-stage pipeline. Stages are numbered sequentially starting from 0, and this ID is used to reference stages in inter-stage dependencies (e.g., `engine_input_source`).
### `runtime`
Configuration for disaggregated execution of the stage, controlling how the stage is deployed and executed.
#### `runtime.process`
Whether to run this stage in a separate process. When set to `true`, the stage will be executed in an isolated process, enabling better resource isolation and parallel execution across different stages. This is essential for multi-GPU deployments where different stages run on different devices.
Default: `true`
#### `runtime.devices`
Visible devices for this stage, specified as a string. This controls which GPU devices are available to the stage process, similar to setting `CUDA_VISIBLE_DEVICES` or using `torch.cuda.set_device()`. For example, `"0"` uses GPU 0, `"1"` uses GPU 1, and `"0,1"` makes both GPUs 0 and 1 visible.
Default: `"0"`
#### `runtime.max_batch_size`
The maximum batch size for offline inference in this stage. This limits how many sequences can be processed together in a single batch during offline inference operations.
Default: `1`
### `engine_args`
Engine arguments for configuring the LLM engine, diffusion engine, or other engine types used by this stage.
#### `engine_args.model_stage`
The name identifier for this model stage within the multi-stage architecture. This is used internally to distinguish different stages of the same model (e.g., "thinker", "talker", "code2wav" in Qwen2.5-Omni).
#### `engine_args.model_arch`
The model architecture class name that is registered in `model_executor/models/registry.py`. This specifies which model implementation to use for this stage. The class must be registered in the model registry for vLLM-Omni to locate and instantiate it.
#### `engine_args.worker_cls`
The specific worker class to use for this stage. This determines how the model computations are executed. Examples include `vllm_omni.worker.gpu_ar_worker.GPUARWorker` for autoregressive stages and `vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker` for diffusion-based stages.
#### `engine_args.scheduler_cls`
The scheduler class to use for this stage. The scheduler manages request queuing, batching, and execution order. Examples include `vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler` for standard stages and `vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler` for diffusion stages.
#### `engine_args.gpu_memory_utilization`
The fraction of GPU memory to allocate for this stage within a single GPU chip. This is a value between 0.0 and 1.0, where 0.8 means 80% of the GPU memory will be used by this stage. This allows fine-grained control over memory allocation when multiple stages share the same GPU or when reserving memory for other operations.
Default: `0.8`
!!! tip "Memory Configuration Guide"
For detailed information on how to calculate memory requirements and properly configure `gpu_memory_utilization`, see the [GPU Memory Calculation and Configuration Guide](./gpu_memory_utilization.md).
#### `engine_args.enforce_eager`
Whether to enforce eager execution mode. When set to `true`, the engine will run in eager mode without using CUDA graphs or other compilation optimizations. Currently, vLLM-Omni only supports eager mode.
Default: `true`
#### `engine_args.trust_remote_code`
Whether to trust remote code when loading models from Hugging Face. This is required for models that use custom code in their configuration files. Set to `true` when loading models that require custom model implementations.
Default: `true`
#### `engine_args.engine_output_type`
Specifies the type of output produced by this stage's engine. This determines what kind of data flows to downstream stages. Possible values include `latent` (hidden states), `text` (tokenized text), and `audio` (audio waveforms). When set to `latent`, the stage outputs latent hidden states in addition to token IDs, which are consumed by downstream stages.
Default: `latent`
#### `engine_args.enable_prefix_caching`
Whether to enable prefix caching for this stage. Prefix caching can improve performance by caching KV cache for common prompt prefixes. However, for requests that output hidden states (when `engine_output_type` is `latent`), prefix caching is not currently supported and should be set to `false`.
Default: `false`
### `is_comprehension`
Whether this stage is a text or multimodal comprehension module. When set to `true`, the stage acts as a comprehension module that processes input text or multimodal content. If this is the first comprehension stage, `AsyncOmni` will use its tokenizer as the default tokenizer for the entire pipeline.
Default: `true`
### `final_output`
Whether this stage produces output that is part of the final outputs returned to the user. When set to `false`, the stage only works as an intermediate stage, processing data that flows to downstream stages but not contributing directly to the final response.
Default: `true`
### `final_output_type`
The type of final output produced by this stage. This specifies what format the output will be in when returned to the user. Currently supported values are `text` (for text generation) and `audio` (for audio generation).
Default: `text`
### `default_sampling_params`
Default sampling parameters for this stage. These parameters control the generation behavior and align with vLLM's sampling parameter semantics. These defaults are used when no explicit sampling parameters are provided in the request.
#### `default_sampling_params.temperature`
Sampling temperature for controlling randomness. Lower values (e.g., 0.0) make the output more deterministic and focused, while higher values increase randomness.
Default: `0.0`
#### `default_sampling_params.top_p`
Nucleus sampling parameter. Only tokens with cumulative probability mass up to `top_p` are considered. This helps filter out low-probability tokens while maintaining diversity.
Default: `1.0`
#### `default_sampling_params.top_k`
Top-k sampling parameter. Only the top `k` most likely tokens are considered. Set to `-1` to disable top-k filtering and consider all tokens.
Default: `-1`
#### `default_sampling_params.max_tokens`
Maximum number of tokens to generate in this stage. This limits the length of the output sequence.
Default: `2048`
#### `default_sampling_params.seed`
Random seed for reproducible generation. When set, the random number generator will be initialized with this seed to ensure consistent outputs across runs.
Default: `42`
#### `default_sampling_params.detokenize`
Whether to detokenize the output tokens into text. When set to `true`, token IDs are converted back to readable text strings.
Default: `True`
#### `default_sampling_params.repetition_penalty`
Penalty applied to tokens that have already appeared in the generated sequence. Values greater than 1.0 discourage repetition, while values less than 1.0 encourage it. A value of 1.0 applies no penalty.
Default: `1.1`
@@ -0,0 +1,94 @@
# stage config for running qwen2.5-omni with architecture of OmniLLM.
stage_args:
- stage_id: 0
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1
engine_args:
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.8
enforce_eager: true # Now we only support eager mode
trust_remote_code: true
engine_output_type: latent
enable_prefix_caching: false
is_comprehension: true
final_output: true
final_output_type: text
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 2048
seed: 42
detokenize: True
repetition_penalty: 1.1
- stage_id: 1
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.8
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
engine_output_type: latent
engine_input_source: [0]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
default_sampling_params:
temperature: 0.9
top_p: 0.8
top_k: 40
max_tokens: 2048
seed: 42
detokenize: True
repetition_penalty: 1.05
stop_token_ids: [8294]
- stage_id: 2
runtime:
process: true
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
max_batch_size: 1
engine_args:
model_stage: code2wav
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
gpu_memory_utilization: 0.15
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
engine_output_type: audio
engine_input_source: [1]
final_output: true
final_output_type: audio
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 2048
seed: 42
detokenize: True
repetition_penalty: 1.1
# Top-level runtime config (concise): default windows and stage edges
runtime:
enabled: true
defaults:
window_size: -1 # Simplified: trigger downstream only after full upstream completion
max_inflight: 1 # Simplified: process serially within each stage
edges:
- from: 0 # thinker → talker: trigger only after receiving full input (-1)
to: 1
window_size: -1
- from: 1 # talker → code2wav: trigger only after receiving full input (-1)
to: 2
window_size: -1
+139
View File
@@ -0,0 +1,139 @@
# Documentation Build Guide
This directory contains the source files for the vLLM-Omni documentation.
## Building Documentation Locally
### Prerequisites
Install documentation dependencies:
```bash
uv pip install -e ".[docs]"
```
### Build and Serve Documentation
From the project root:
```bash
# Serve documentation locally (auto-reload on changes)
# This starts a local web server at http://127.0.0.1:8000
mkdocs serve
# Build static site (generates HTML files in site/ directory)
mkdocs build
```
When using `mkdocs serve`, the documentation will be automatically available at `http://127.0.0.1:8000`. The server will automatically reload when you make changes to the documentation files.
## Auto-generating API Documentation
The documentation automatically extracts docstrings from the code using mkdocstrings. To ensure your code is documented:
1. Add docstrings to all public classes, functions, and methods
2. Use Google or NumPy style docstrings (both are supported)
3. Rebuild the documentation to see changes
Example docstring:
```python
class Omni:
"""Main entry point for vLLM-Omni inference.
This class provides a high-level interface for running multi-modal
inference with non-autoregressive models.
Args:
model: Model name or path
stage_configs: Optional stage configurations
**kwargs: Additional arguments passed to the engine
Example:
>>> llm = Omni(model="Qwen/Qwen2.5-Omni")
>>> outputs = llm.generate(prompts="Hello")
"""
```
## Documentation Structure
```
docs/
├── index.md # Main documentation page
├── getting_started/ # Getting started guides
├── architecture/ # Architecture documentation
├── api/ # API reference (auto-generated from code)
├── examples/ # Code examples
└── stylesheets/ # Custom CSS
```
## Publishing Documentation
### GitHub Pages (Recommended)
The documentation is automatically deployed to GitHub Pages using GitHub Actions.
1. **Enable GitHub Pages**:
- Go to repository `Settings` → `Pages`
- Set `Source` to `GitHub Actions`
- Save settings
2. **Push changes**:
```bash
git push origin main
```
3. **Documentation will be available at**:
- `https://vllm-omni.readthedocs.io`
The GitHub Actions workflow (`.github/workflows/docs.yml`) will automatically:
- Build the documentation when you push to `main` branch
- Deploy it to GitHub Pages
- Update the documentation whenever you make changes
### Read the Docs (Alternative)
You can also use Read the Docs for hosting:
1. Sign up at https://readthedocs.org/
2. Import the `vllm-project/vllm-omni` repository
3. Read the Docs will automatically build using `.readthedocs.yml`
4. Documentation will be available at: `https://vllm-omni.readthedocs.io/`
## Configuration
The documentation configuration is in `mkdocs.yml` at the project root.
## Tips
- **API Documentation**: API docs are automatically generated using `mkdocs-api-autonav` and `mkdocstrings`
- No need to manually create API pages - they're generated automatically
- Use `[module.name.ClassName][]` syntax for cross-references in Summary pages
- **Code Snippets**: Use `--8<-- "path/to/file.py"` for including code snippets
- **Markdown**: Use Markdown for all documentation (no need for RST)
- **Material Theme**: Use Material theme features like:
- Admonitions: `!!! note`, `!!! warning`, etc.
- Code blocks with syntax highlighting
- Tabs for organizing content
- Math formulas using `pymdownx.arithmatex`
## Troubleshooting
### Documentation not updating
- Make sure you've saved all files
- If using `mkdocs serve`, it should auto-reload
- Check for syntax errors in `mkdocs.yml`
### API links not working
- Ensure class names match exactly (case-sensitive)
- Check that the module is imported correctly
- Run `mkdocs build --strict` to check for errors
### Build errors
- Check Python version (requires 3.9+)
- Ensure all dependencies are installed: `pip install -e ".[docs]"`
- Check `mkdocs.yml` syntax with `mkdocs build --strict`
+155
View File
@@ -0,0 +1,155 @@
# Contributing to vLLM-Omni
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).
## Getting Started
vLLM-Omni uses `uv` as the environment manager, to create and manage Python environments. Please follow the documentation to install `uv`. After installing `uv`, you can create a new Python environment using the following commands:
```bash
uv venv --python 3.12 --seed
source .venv/bin/activate
```
### Development Environment for vLLM and vLLM-Omni
vLLM-Omni is quickly evolving, please see the [installation guide](../getting_started/installation/README.md) for details. It's recommended to build from source to provide the latest development environment.
!!! tip
vLLM-Omni is compatible with Python versions 3.10 to 3.12. However, we recommend developing with Python 3.12 to minimize the chance of your local environment clashing with our CI environment.
### Adding a new model to vLLM-Omni
Please check [model implementation](model/README.md) for how to add diffusion and omni-modality models to vLLM-Omni.
### Linting
vLLM-Omni uses `pre-commit` to lint and format the codebase. See [pre-commit documentation](https://pre-commit.com/#usage) if `pre-commit` is new to you. Setting up `pre-commit` is as easy as:
```bash
uv pip install pre-commit
pre-commit install
```
vLLM-Omni's `pre-commit` hooks will now run automatically every time you commit.
!!! tip
You can manually run the `pre-commit` hooks using:
```bash
pre-commit run # runs on staged files
pre-commit run --show-diff-on-failure --color=always --all-files # runs on all files (short for --all-files)
```
### Documentation
MkDocs is a fast, simple and downright gorgeous static site generator that's geared towards building project documentation. Documentation source files are written in Markdown, and configured with a single YAML configuration file, `mkdocs.yml`.
Get started with:
```bash
uv pip install -e ".[docs]"
```
MkDocs comes with a built-in dev-server that lets you preview your documentation as you work on it. From the root of the repository, run:
```bash
mkdocs serve # with API ref (~10 minutes)
API_AUTONAV_EXCLUDE=vllm_omni mkdocs serve # API ref off (~15 seconds)
```
Once you see `Serving on http://127.0.0.1:8000/` in the logs, the live preview is ready! Open <http://127.0.0.1:8000/> in your browser to see it.
For additional features and advanced configurations, refer to the:
- [MkDocs documentation](https://www.mkdocs.org/)
- [Material for MkDocs documentation](https://squidfunk.github.io/mkdocs-material/) (the MkDocs theme we use)
### Testing
vLLM-Omni uses `pytest` to test the codebase.
```bash
# Run all tests
pytest tests/
# Run tests for a single test file with detailed output
pytest -s -v tests/test_omni_llm.py
```
!!! warning
Currently, not all unit tests pass when run on CPU platforms. If you don't have access to a GPU platform to run unit tests locally, rely on the continuous integration system to run the tests for now.
## Issues
If you encounter a bug or have a feature request, please search existing issues first to see if it has already been reported. If not, please file a new issue, providing as much relevant information as possible.
!!! important
If you discover a security vulnerability, please report it by creating a GitHub issue with the `security` label.
## Pull Requests & Code Reviews
Thank you for your contribution to vLLM-Omni! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM-Omni maintain the code quality and improve the efficiency of the review process.
### DCO and Signed-off-by
When contributing changes to this project, you must agree to the [DCO](https://developercertificate.org/). Commits must include a `Signed-off-by:` header which certifies agreement with the terms of the DCO.
Using `-s` with `git commit` will automatically add this header.
!!! tip
You can enable automatic sign-off via your IDE:
- **PyCharm**: Click on the `Show Commit Options` icon to the right of the `Commit and Push...` button in the `Commit` window. It will bring up a `git` window where you can modify the `Author` and enable `Sign-off commit`.
- **VSCode**: Open the Settings editor and enable the `Git: Always Sign Off` (`git.alwaysSignOff`) field.
### PR Title and Classification
Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:
- `[Bugfix]` for bug fixes.
- `[CI/Build]` for build or continuous integration improvements.
- `[Doc]` for documentation fixes and improvements.
- `[Model]` for adding a new model or improving an existing model. Model name should appear in the title.
- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `OmniLLM` class, etc.)
- `[Kernel]` for changes affecting CUDA kernels or other compute kernels.
- `[Core]` for changes in the core vLLM-Omni logic (e.g., `OmniProcessor`, `OmniARScheduler`, etc.)
- `[Hardware][Vendor]` for hardware-specific changes. Vendor name should appear in the prefix, such as [Ascend] for Ascend NPUs.
- `[Misc]` for PRs that do not fit the above categories. Please use this sparingly.
!!! note
If the PR spans more than one category, please include all relevant prefixes.
### Code Quality
The PR needs to meet the following code quality standards:
- We adhere to Google Python style guide and Google C++ style guide.
- Pass all linter checks.
- The code needs to be well-documented to ensure future contributors can easily understand the code.
- Include sufficient tests to ensure the project stays correct and robust. This includes both unit tests and integration tests.
- Please add documentation to `docs/` if the PR modifies the user-facing behaviors of vLLM-Omni. It helps vLLM-Omni users understand and utilize the new features or changes.
### Notes for Large Changes
Please keep the changes as concise as possible. For major architectural changes (>500 LOC excluding kernel/data/config/test), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with `rfc-required` and might not go through the PR.
### What to Expect for the Reviews
The goal of the vLLM-Omni team is to be a _transparent reviewing machine_. We would like to make the review process transparent and efficient and make sure no contributor feels confused or frustrated. However, the vLLM-Omni team is small, so we need to prioritize some PRs over others. Here is what you can expect from the review process:
- After the PR is submitted, the PR will be assigned to a reviewer. Every reviewer will pick up the PRs based on their expertise and availability.
- After the PR is assigned, the reviewer will provide status updates every 2-3 days. If the PR is not reviewed within 7 days, please feel free to ping the reviewer or the vLLM-Omni team.
- After the review, the reviewer will put an `action-required` label on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR.
- Please respond to all comments within a reasonable time frame. If a comment isn't clear or you disagree with a suggestion, feel free to ask for clarification or discuss the suggestion.
## Additional Resources
- [Design Documents](../design/index.md) - Architecture and design documentation
## Thank You
Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM-Omni. All of your contributions help make vLLM-Omni a great tool and community for everyone!
+4
View File
@@ -0,0 +1,4 @@
# CI Failures
What should I do when a CI job fails on my PR, but I don't think my PR caused
the failure?
+160
View File
@@ -0,0 +1,160 @@
# Markers for Tests
By adding markers before test functions, tests can later be executed uniformly by simply declaring the corresponding marker type.
## Current Markers
Defined in `pyproject.toml`:
| Marker | Description |
| ------------------ | ------------------------------------------------------- |
| `core_model` | Core model tests (run in each PR) |
| `diffusion` | Diffusion model tests |
| `omni` | Omni model tests |
| `cache` | Cache backend tests |
| `parallel` | Parallelism/distributed tests |
| `cpu` | Tests that run on CPU |
| `gpu` | Tests that run on GPU (auto-added) |
| `cuda` | Tests that run on CUDA (auto-added) |
| `rocm` | Tests that run on AMD/ROCm (auto-added) |
| `npu` | Tests that run on NPU/Ascend (auto-added) |
| `H100` | Tests that require H100 GPU |
| `L4` | Tests that require L4 GPU |
| `MI325` | Tests that require MI325 GPU (AMD/ROCm) |
| `A2` | Tests that require A2 NPU |
| `A3` | Tests that require A3 NPU |
| `distributed_cuda` | Tests that require multi cards on CUDA platform |
| `distributed_rocm` | Tests that require multi cards on ROCm platform |
| `distributed_npu` | Tests that require multi cards on NPU platform |
| `skipif_cuda` | Skip if the num of CUDA cards is less than the required |
| `skipif_rocm` | Skip if the num of ROCm cards is less than the required |
| `skipif_npu` | Skip if the num of NPU cards is less than the required |
| `slow` | Slow tests (may skip in quick CI) |
| `benchmark` | Benchmark tests |
For those markers shown as auto-added, they will be added by the `@hardware_test` decorator.
### Example usage for markers
```python
from tests.utils import hardware_test
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
num_cards=2,
)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_video_to_audio()
...
```
### Decorator: `@hardware_test`
This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/utils.py` performs the following actions:
1. **Applies platform and resource markers**
Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `A2`, `A3`).
```
@pytest.mark.cuda
@pytest.mark.L4
```
2. **Handles multi-card (distributed) scenarios**
For tests requiring multiple cards, it automatically adds distributed markers such as `distributed_cuda`, `distributed_rocm`, or `distributed_npu`.
```
@pytest.mark.distributed_cuda(num_cards=num_cards)
```
3. **Supports flexible card requirements**
Accepts `num_cards` as either a single integer for all platforms or as a dictionary with per-platform values. If not specified, defaults to 1 card per platform.
4. **Integrates resource validation**
On CUDA, adds a skip marker (`skipif_cuda`) if the system does not have the required number of devices.
Support for `skipif_rocm` and `skipif_npu` will be implemented later.
5. **Runs each test in a new process**
Automatically wraps the distributed test with a decorator (`@create_new_process_for_each_test`) to ensure isolation and compatibility with multi-process hardware backends.
6. **Works with pytest filtering**
Allows tests to be filtered and selected at runtime using standard pytest marker expressions (e.g., `-m "distributed_cuda and L4"`).
#### Example usage for decorator
- Single call for multiple platforms:
```python
@hardware_test(
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
num_cards={"cuda": 2, "rocm": 2, "npu": 2},
)
```
or
```python
@hardware_test(
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
num_cards=2,
)
```
- `res` must be a dict; supported resources: CUDA (L4/H100), ROCm (MI325), NPU (A2/A3)
- `num_cards` can be int (all platforms) or dict (per platform); defaults to 1 when missing
- `hardware_test` automatically applies `@create_new_process_for_each_test` for distributed tests.
- Distributed markers (`distributed_cuda`, `distributed_rocm`, `distributed_npu`) are auto-added for multi-card cases
- Filtering examples:
- CUDA only: `pytest -m "distributed_cuda and L4"`
- ROCm only: `pytest -m "distributed_rocm and MI325"`
- NPU only: `pytest -m "distributed_npu"`
## Add Support for a New Platform
If you want to add support for a new platform (e.g., "tpu" for a new accelerator), follow these steps:
1. **Extend the marker list in your pytest config** so that platform/resource markers are defined:
```toml
# In pyproject.toml or pytest.ini
[tool.pytest.ini_options]
markers = [
# ... existing markers ...
"tpu: Tests that require TPU device",
"TPU_V3: Tests that require TPU v3 hardware",
"distributed_tpu: Tests that require multiple TPU devices",
]
```
2. **Implement a marker construction function for your platform** in `vllm-omni/tests/utils.py`:
```python
# In vllm-omni/tests/utils.py
def tpu_marks(*, res: str, num_cards: int):
test_platform = pytest.mark.tpu
if res == "TPU_V3":
test_resource = pytest.mark.TPU_V3
else:
raise ValueError(
f"Invalid TPU resource type: {res}. Supported: TPU_V3")
if num_cards == 1:
return [test_platform, test_resource]
else:
test_distributed = pytest.mark.distributed_tpu(num_cards=num_cards)
# Optionally: add skipif_tpu when implemented
return [test_platform, test_resource, test_distributed]
```
3. **Update `hardware_test` to recognize your new platform**:
In the relevant place (see the `hardware_test` implementation), add:
```python
if platform == "tpu":
marks = tpu_marks(res=resource, num_cards=cards)
```
4. **(Recommended) Add a test using your new markers**:
```python
@hardware_test(
res={"tpu": "TPU_V3"},
num_cards=2,
)
def test_my_tpu_feature():
...
```
**Summary**:
- Add pytest markers for your new platform/resources
- Implement a marker function (`xxx_marks`)
- Plug into `hardware_test`
- You're done: tests decorated with `@hardware_test` using your platform now automatically get the correct markers, distribution, and isolation!
See code in `vllm-omni/tests/utils.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).
+280
View File
@@ -0,0 +1,280 @@
# Test File Structure and Style Guide
To ensure project maintainability and sustainable development, we encourage contributors to submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes. This document outlines the guidelines for organizing and naming test files.
## Test Types
### Unit Tests and System Tests
For unit tests and system tests, we strongly recommend placing test files in the same directory structure as the source code being tested, using the naming convention `test_*.py`.
### End-to-End (E2E) Tests for Models
End-to-end tests verify the complete functionality of a system or component. For our project, the E2E tests for different omni models are organized into two subdirectories:
- **`tests/e2e/offline_inference/`**: Tests for offline inference modes (e.g., Qwen3Omni offline inference)
- **`tests/e2e/online_serving/`**: Tests for online serving scenarios (e.g., API server tests)
**Example:** The test file for `vllm_omni/entrypoints/omni_llm.py` should be located at `tests/entrypoints/test_omni_llm.py`.
## Test Directory Structure
The ideal directory structure mirrors the source code organization:
```
vllm_omni/ tests/
├── config/ → ├── config/
│ └── model.py │ └── test_model.py
│
├── core/ → ├── core/
│ └── sched/ │ └── sched/ # Maps to core/sched/
│ ├── omni_ar_scheduler.py │ ├── test_omni_ar_scheduler.py
│ ├── omni_generation_scheduler.py │ ├── test_omni_generation_scheduler.py
│ └── output.py │ └── test_output.py
│
├── diffusion/ → ├── diffusion/
│ ├── diffusion_engine.py │ ├── test_diffusion_engine.py
│ ├── omni_diffusion.py │ ├── test_omni_diffusion.py
│ ├── attention/ │ ├── attention/ # Maps to diffusion/attention/
│ │ └── backends/ │ │ └── test_*.py
│ ├── models/ │ ├── models/ # Maps to diffusion/models/
│ │ ├── qwen_image/ │ │ ├── qwen_image/
│ │ │ └── ... │ │ │ └── test_*.py
│ │ └── z_image/ │ │ └── z_image/
│ │ └── ... │ │ └── test_*.py
│ └── worker/ │ └── worker/ # Maps to diffusion/worker/
│ └── ... │ └── test_*.py
│
├── distributed/ → ├── distributed/
│ └── ... │ └── test_*.py
│
├── engine/ → ├── engine/
│ ├── processor.py │ ├── test_processor.py
│ └── output_processor.py │ └── test_output_processor.py
│
├── entrypoints/ → ├── entrypoints/
│ ├── omni_llm.py │ ├── test_omni_llm.py # UT: OmniLLM core logic (mocked)
│ ├── omni_stage.py │ ├── test_omni_stage.py # UT: OmniStage logic
│ ├── omni.py │ ├── test_omni.py # E2E: Omni class (offline inference)
│ ├── async_omni.py │ ├── test_async_omni.py # E2E: AsyncOmni class
│ ├── cli/ │ ├── cli/ # Maps to entrypoints/cli/
│ │ └── ... │ │ └── test_*.py
│ └── openai/ │ └── openai/ # Maps to entrypoints/openai/
│ ├── api_server.py │ ├── test_api_server.py # E2E: API server (online serving)
│ └── serving_chat.py │ └── test_serving_chat.py
│
├── inputs/ → ├── inputs/
│ ├── data.py │ ├── test_data.py
│ ├── parse.py │ ├── test_parse.py
│ └── preprocess.py │ └── test_preprocess.py
│
├── model_executor/ → ├── model_executor/
│ ├── layers/ │ ├── layers/
│ │ └── mrope.py │ │ └── test_mrope.py
│ ├── model_loader/ │ ├── model_loader/
│ │ └── weight_utils.py │ │ └── test_weight_utils.py
│ ├── models/ │ ├── models/
│ │ ├── qwen2_5_omni/ │ │ ├── qwen2_5_omni/
│ │ │ ├── qwen2_5_omni_thinker.py │ │ │ ├── test_qwen2_5_omni_thinker.py # UT
│ │ │ ├── qwen2_5_omni_talker.py │ │ │ ├── test_qwen2_5_omni_talker.py # UT
│ │ │ └── qwen2_5_omni_token2wav.py │ │ │ └── test_qwen2_5_omni_token2wav.py # UT
│ │ └── qwen3_omni/ │ │ └── qwen3_omni/
│ │ └── ... │ │ └── test_*.py
│ ├── stage_configs/ │ └── stage_configs/ # Configuration tests (if needed)
│ │ └── ... │ └── test_*.py
│ └── stage_input_processors/ │ └── stage_input_processors/
│ └── ... │ └── test_*.py
│
├── sample/ → ├── sample/
│ └── ... │ └── test_*.py
│
├── utils/ → ├── utils/
│ └── platform_utils.py │ └── test_platform_utils.py
│
├── worker/ → ├── worker/
├── gpu_ar_worker.py │ ├── test_gpu_ar_worker.py
├── gpu_generation_worker.py │ ├── test_gpu_generation_worker.py
├── gpu_model_runner.py │ ├── test_gpu_model_runner.py
└── npu/ │ └── npu/ # Maps to worker/npu/
└── ... │ └── test_*.py
│
└── e2e/ → ├── e2e/ # End-to-end scenarios (no 1:1 source mirror)
├── online_serving/ # Full-stack online serving flows
│ └── (empty for now)
└── offline_inference/ # Full offline inference flows
├── test_qwen2_5_omni.py # Moved from multi_stages/
├── test_qwen3_omni.py # Moved from multi_stages_h100/
├── test_t2i_model.py # Moved from single_stage/
└── stage_configs/ # Shared stage configs
├── qwen2_5_omni_ci.yaml
└── qwen3_omni_ci.yaml
```
### Naming Conventions
- **Unit Tests**: Use `test_<module_name>.py` format. Example: `omni_llm.py` → `test_omni_llm.py`
- **E2E Tests**: Place in `tests/e2e/offline_inference/` or `tests/e2e/online_serving/` with descriptive names. Example: `tests/e2e/offline_inference/test_qwen3_omni.py`, `tests/e2e/offline_inference/test_diffusion_model.py`
### Best Practices
1. **Mirror Source Structure**: Test directories should mirror the source code structure
2. **Test Type Indicators**: Use comments to indicate test types (UT for unit tests, E2E for end-to-end tests)
3. **Shared Resources**: Place shared test configurations (e.g., CI configs) in appropriate subdirectories
4. **Consistent Naming**: Follow the `test_*.py` naming convention consistently across all test files
## Test codes requirements
### Coding style
1. **File header**: Add SPDX license header to all test files
2. **Imports**: Pls don't use manual `sys.path` modifications, use standard imports instead.
3. **Test type differentiation**:
- Unit tests: Maintain mock style
- E2E tests for models: Consider using OmniRunner uniformly, avoid decorators
4. **Documentation**: Add docstrings to all test functions
5. **Environment variables**: Set uniformly in `conftest.py` or at the top of files
6. **Type annotations**: Add type annotations to all test function parameters
7. **Pytest Markers**: Add necessary markers like `@pytest.mark.core_model` and use `@hardware_test` to declare hardware requirements (check detailed in [Markers for Tests](../ci/tests_markers.md)).
### Template
#### E2E - Online serving
```python
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Online E2E smoke test for an omni model (video,text,audio → audio).
"""
from pathlib import Path
import pytest
import openai
from tests.utils import hardware_test
# Optional: set process start method for workers
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
models = ["{your model name}"] #Edit here to load your model
stage_configs = [str(Path(__file__).parent / "stage_configs" / {your model yaml})] #Edit here to load your model yaml
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
#OmniServer,Used to start the vllm-omni server
class OmniServer:
xxx
@pytest.fixture
def omni_server(request):
model, stage_config_path = request.param
with OmniServer(model, ["--stage-configs-path", stage_config_path]) as server:
yield server
#handle request message
@pytest.fixture(scope="session")
def base64_encoded_video() -> str:
xxx
@pytest.fixture(scope="session")
def dummy_messages_from_video_data(video_data_url: str, content_text: str) -> str:
xxx
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
num_cards={"cuda": 2, "rocm": 2, "npu": 4},
)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_video_to_audio(
client: openai.OpenAI,
omni_server,
base64_encoded_video: str,
) -> None:
#set message
video_data_url = f"data:video/mp4;base64, {base64_encoded_video}"
messages = dummy_messages_from_video_data(video_data_url)
#send request
chat_completion = client.chat.completions.create(
model=omni_server.model,
messages=messages,
)
#verify text output
text_choice = chat_completion.choices[0]
assert text_choice.finish_reason == "length"
#verify audio output
audio_choice = chat_completion.choices[1]
audio_message = audio_choice.message
if hasattr(audio_message, "audio") and audio_message.audio:
assert audio_message.audio.data is not None
assert len(audio_message.audio.data) > 0
```
#### E2E - Offline inference
```python
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Offline E2E smoke test for an omni model (video → audio).
"""
import os
from pathlib import Path
import pytest
from vllm.assets.video import VideoAsset
from tests.utils import hardware_test
from ..multi_stages.conftest import OmniRunner
# Optional: set process start method for workers
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
models = ["{your model name}"] #Edit here to load your model
stage_configs = [str(Path(__file__).parent / "stage_configs" / {your model yaml})] #Edit here to load your model yaml
# Create parameter combinations for model and stage config
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
# function name: test_{input_modality}_to_{output_modality}
# modality candidate: text, image, audio, video, mixed_modalities
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
num_cards=2,
)
@pytest.mark.parametrize("test_config", test_params)
def test_video_to_audio(omni_runner: type[OmniRunner], model: str) -> None:
"""Offline inference: video input, audio output."""
model, stage_config_path = test_config
with omni_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
# Prepare inputs
video = VideoAsset(name="sample", num_frames=4).np_ndarrays
outputs = runner.generate_multimodal(
prompts="Describe this video briefly.",
videos=video,
)
# Minimal assertions: got outputs and at least one audio result
assert outputs
has_audio = any(o.final_output_type == "audio" for o in outputs)
assert has_audio
```
## Checklist before submitting your test files
1. The file is saved in an appropriate place and the file name is clear.
2. The coding style follows the requirements outlined above.
3. **All test functions have appropriate pytest markers**
4. For tests that need run in CI, please ensure the test is configured under the `./buildkite/` folder.
+15
View File
@@ -0,0 +1,15 @@
# Adding a New Model
This section provides comprehensive guidance on how to add a new model to vLLM-Omni.
## Documentation
- **[Adding an Omni-Modality Model](adding_omni_model.md)**: Complete step-by-step guide using Qwen3-Omni as an example.
- **[Adding a Diffusion Model](adding_diffusion_model.md)**: Complete step-by-step guide using Qwen/Qwen-Image-Edit as an example.
## Quick Start
For a quick reference, see the [Adding a New Multi-Stage Model Guide](adding_omni_model.md) and [Adding a New Diffusion Model Guide](adding_diffusion_model.md).
@@ -0,0 +1,227 @@
# Adding a Diffusion Model
This guide walks through the process of adding a new Diffusion model to vLLM-Omni, using Qwen/Qwen-Image-Edit as a comprehensive example.
# Table of Contents
1. [Overview](#overview)
2. [Directory Structure](#directory-structure)
3. [Step-by-Step Implementation](#step-by-step-implementation)
4. [Testing](#testing)
5. [Adding a Model Recipe](#adding-a-model-recipe)
# Overview
When add a new diffusion model into vLLM-Omni, additional adaptation work is required due to the following reasons:
+ New model must follow the framework’s parameter passing mechanisms and inference flow.
+ Replacing the model’s default implementations with optimized modules, which is necessary to achieve the better performance.
The diffusion execution flow as follow:
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-diffusion-flow.png">
<img alt="Diffusion Flow" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-diffusion-flow.png" width=55%>
</picture>
</p>
# Directory Structure
File Structure for Adding a New Diffusion Model
```
vllm_omni/
└── examples/
└──offline_inference
└── example script # reuse existing if possible (e.g., image_edit.py)
└──online_serving
└── example script
└── diffusion/
└── registry.py # Registry work
├── request.py # Request Info
└── models/your_model_name/ # Model directory (e.g., qwen_image)
└── pipeline_xxx.py # Model implementation (e.g., pipeline_qwen_image_edit.py)
```
# Step-by-step-implementation
## Step 1: Model Implementation
The diffusion pipeline’s implementation follows **HuggingFace Diffusers**.
### 1.1 Define the Pipeline Class
Define the pipeline class, e.g., `QwenImageEditPipeline`, and initialize all required submodules, either from HuggingFace `diffusers` or custom implementations. In `QwenImageEditPipeline`, only `QwenImageTransformer2DModel` is re-implemented to support optimizations such as Ulysses-SP. When adding new models in the future, you can either reuse this re-implemented `QwenImageTransformer2DModel` or extend it as needed.
### 1.2 Pre-Processing and Post-Processing Extraction
Extract the pre-processing and post-processing logic from the pipeline class to follow vLLM-Omni’s execution flow. For Qwen-Image-Edit:
```python
def get_qwen_image_edit_pre_process_func(
od_config: OmniDiffusionConfig,
):
"""
Define a pre-processing function that resizes input images and
pre-process for subsequent inference.
"""
```
```python
def get_qwen_image_edit_post_process_func(
od_config: OmniDiffusionConfig,
):
"""
Defines a post-processing function that post-process images.
"""
```
### 1.3 Define the forward function
The forward function of `QwenImageEditPipeline` follows the HuggingFace `diffusers` design for the most part. The key differences are:
+ As described in the overview, arguments are passed through `OnniDiffusionRequest`, so we need to get user parameters from it accordingly.
```python
prompt = req.prompt
```
+ pre/post-processing are handled by the framework elsewhere, so skip them.
### 1.4 Replace some ops or layers in DiT component
vLLM-Omni provides a set of optimized operators with better performance and built-in support for parallelism, including attention, rotary embeddings (RoPE), and linear layers.
Below is an example showing how to replace standard Transformer attention and FFN layers with vLLM-Omni implementations:
```python
from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
class MyAttention(nn.Module):
def __init__(self):
super().__init__()
self.attn = Attention()
self.to_qkv = QKVParallelLinear()
self.to_out = RowParallelLinear()
self.rope = RotaryEmbedding(is_neox_style=False)
def forward(self, hidden_states):
qkv, _ = self.to_qkv(hidden_states)
q, k, v = qkv.split(...)
q, k = self.rope(...)
attn_output = self.attn(q, k, v)
output = self.to_out(attn_output)
class MyFFN(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = ColumnParallelLinear()
self.fc2 = RowParallelLinear()
self.act = F.gelu
def forward(self, hidden_states):
hidden, _ = self.fc1(hidden_states)
hidden = self.act(hidden)
output = self.fc2(hidden)
return output
```
In this example:
+ Attention uses vLLM-Omni’s optimized attention kernel together with parallel QKV projection and RoPE.
+ Linear layers are replaced with column- and row-parallel variants to enable tensor parallelism.
+ The FFN follows a standard two-layer structure and can be further optimized (e.g., using fused or merged projections) if needed.
### 1.5 Provide a `_repeated_blocks` in DiT model
`_repeated_blocks` is the small and frequently-repeated block(s) of a model -- typically a transformer layer.
It's used for torch compile optimizations.
```python
_repeated_blocks = ["QwenImageTransformerBlock"]
```
### 1.6 (Optional) implement sequence parallelism
vLLM-Omni has a non-intrusive `_sp_plan` that enable sequence parallel without modifying `forward()` logic.
You can refer to [How to parallelize a new model](../../user_guide/diffusion/parallelism_acceleration.md)
### 1.7 (Optional) integrate with Cache-Dit
vLLM-Omni supports acceleration via [Cache-Dit](../../user_guide/diffusion/cache_dit_acceleration.md). Most models compatible with Diffusers can use Cache-Dit seamlessly. For new models, you can extend support by modifying`cache_dit_backend.py`
## Step 2: Extend OmniDiffusionRequest Fields
User-provided inputs are ultimately passed to the model’s forward method through OmniDiffusionRequest, so we add the required fields here to support the new model.
```python
prompt: str | list[str] | None = None
negative_prompt: str | list[str] | None = None
...
```
## Step 3: Registry
+ registry diffusion model in registry.py
```python
_DIFFUSION_MODELS = {
# arch:(mod_folder, mod_relname, cls_name)
...
"QwenImageEditPipeline": (
"qwen_image",
"pipeline_qwen_image_edit",
"QwenImageEditPipeline",
),
...
}
```
+ registry pre-process get function
```python
_DIFFUSION_PRE_PROCESS_FUNCS = {
# arch: pre_process_func
...
"QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func",
...
}
```
+ registry post-process get function
```python
_DIFFUSION_POST_PROCESS_FUNCS = {
# arch: post_process_func
...
"QwenImageEditPipeline": "get_qwen_image_edit_post_process_func",
...
}
```
## Step 4: Add an Example Script
For each newly integrated model, we need to provide examples script under the examples/ to demonstrate how to initialize the pipeline with Omni, pass in user inputs, and generate outputs.
Key point for writing the example:
+ Use the Omni entrypoint to load the model and construct the pipeline.
+ Show how to format user inputs and pass them via omni.generate(...).
+ Demonstrate the common runtime arguments, such as:
+ model path or model name
+ input image(s) or prompt text
+ key diffusion parameters (e.g., inference steps, guidance scale)
+ optional acceleration backends (e.g., Cache-DiT, TeaCache)
+ Save or display the generated results so users can validate the integration.
## Step 5: Open a Pull Request
When submitting a pull request to add support for a new model, please include the following information in the PR description:
+ Output verification: provide generation outputs to verify correctness and model behavior.
+ Inference speed: provide a comparison with the corresponding implementation in Diffusers.
+ Parallelism support: specify the supported parallel sizes and any relevant limitations.
+ Cache acceleration: check whether the model can be accelerated using Cache-Dit or not.
Providing these details helps reviewers evaluate correctness, performance improvements, and parallel scalability of the new model integration.
# Testing
For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
## Adding a Model Recipe
After implementing and testing your model, please add a model recipe to the [vllm-project/recipes](https://github.com/vllm-project/recipes) repository. This helps other users understand how to use your model with vLLM-Omni.
@@ -0,0 +1,633 @@
# Adding an Omni-Modality Model
This guide walks through the process of adding a new multi-stage model to vLLM-Omni, using **Qwen3-Omni** as a comprehensive example. Qwen3-Omni is a multi-stage omni-modality model that demonstrates the full capabilities of vLLM-Omni's architecture.
## Table of Contents
1. [Overview](#overview)
2. [Directory Structure](#directory-structure)
3. [Step-by-Step Implementation](#step-by-step-implementation)
4. [Key Components](#key-components)
5. [Model Registration](#model-registration)
6. [Stage Configuration](#stage-configuration)
7. [Stage Input Processors](#stage-input-processors)
8. [Testing](#testing)
9. [Adding a Model Recipe](#adding-a-model-recipe)
10. [Summary](#summary)
## Overview
vLLM-Omni supports multi-stage model architectures where different stages can run on different devices and process different modalities. The Qwen3-Omni model exemplifies this with three stages:
1. **Thinker Stage**: Multimodal understanding (text + audio + video) → text generation
2. **Talker Stage**: Text embeddings → RVQ codec codes
3. **Code2Wav Stage**: RVQ codes → audio waveform
Each stage is implemented as a separate model class that can be configured independently.
## Directory Structure
When adding a new model, you'll need to create the following structure:
```
vllm_omni/model_executor/models/
└── your_model_name/ # Model directory (e.g., qwen3_omni)
├── __init__.py # Exports main model class
├── your_model.py # Main unified model class
├── your_model_stage1_implementation.py # Stage 1 implementation (e.g., thinker)
├── your_model_stage2_implementation.py # Stage 2 implementation (e.g., talker)
└── your_model_stage3_implementation.py # Stage 3 implementation (e.g., code2wav)
└── ... maybe other stage implementations
vllm_omni/model_executor/stage_input_processors/
└── your_model_name.py # Stage transition processors
vllm_omni/model_executor/stage_configs/
└── your_model_name.yaml # Stage configuration file
```
## Step-by-Step Implementation
### Step 1: Create the Model Directory
Create a new directory under `vllm_omni/model_executor/models/`
### Step 2: Implement Stage Components
For Qwen3-Omni, we have three stage components:
#### 2.1 Thinker Stage (`qwen3_omni_moe_thinker.py`)
The thinker stage handles multimodal understanding. Key features:
- Inherits from base Qwen3 MoE model in vLLM, using vLLM fused ops & page attn to accelerate
- Implements multimodal processing interfaces
- Handles audio, video, and image inputs
- Generates text outputs
```python
from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM
class Qwen3OmniMoeThinkerForConditionalGeneration(
Qwen3MoeForCausalLM,
SupportsMultiModal,
SupportsPP
):
"""Thinker stage: multimodal understanding → text generation."""
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
# Initialize base model
# Set up multimodal processors
# Configure audio/video/image encoders
pass
```
#### 2.2 Talker Stage (`qwen3_omni_moe_talker.py`)
The talker stage converts text embeddings to codec codes:
```python
class Qwen3OmniMoeTalkerForConditionalGeneration(
Qwen3MoeForCausalLM,
SupportsPP
):
"""Talker stage: text embeddings → RVQ codec codes."""
def __init__(self, vllm_config, talker_config, prefix):
# Initialize base model
# Replace LM head with codec head
# Set up text projection from thinker
pass
```
#### 2.3 Code2Wav Stage (`qwen3_omni_code2wav.py`)
The code2wav stage generates audio waveforms:
```python
class Qwen3OmniMoeCode2Wav(nn.Module):
"""Code2Wav stage: RVQ codes → audio waveform."""
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
# Initialize audio decoder
# Set up codec processing
pass
```
### Step 3: Implement the Unified Model Class
The main model class (`qwen3_omni.py`) orchestrates all stages:
```python
@MULTIMODAL_REGISTRY.register_processor(
Qwen3OmniMoeThinkerMultiModalProcessor,
info=Qwen3OmniMoeThinkerProcessingInfo,
dummy_inputs=Qwen3OmniMoeThinkerDummyInputsBuilder,
)
class Qwen3OmniMoeForConditionalGeneration(
nn.Module, SupportsMultiModal, SupportsPP, Qwen3OmniMoeConditionalGenerationMixin
):
"""
Unified Qwen3 Omni MoE model combining thinker, talker, and code2wav.
Architecture:
- Thinker: Multimodal understanding (text + audio + video) → text generation
- Talker: Text embeddings → RVQ codec codes
- Code2Wav: RVQ codes → audio waveform
Usage:
Set `model_stage` in vllm_config to one of: "thinker", "talker", "code2wav"
"""
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
self.have_multimodal_outputs = True
config: Qwen3OmniMoeConfig = vllm_config.model_config.hf_config
# Determine which stage to initialize
self.model_stage = vllm_config.model_config.model_stage
if self.model_stage == "thinker":
# Initialize thinker model
thinker_vllm_config = vllm_config.with_hf_config(
config.thinker_config,
architectures=["Qwen3OmniMoeThinkerForConditionalGeneration"]
)
self.thinker = init_vllm_registered_model(
vllm_config=thinker_vllm_config,
prefix=maybe_prefix(prefix, "thinker"),
hf_config=config.thinker_config,
architectures=["Qwen3OmniMoeThinkerForConditionalGeneration"],
)
self.model = self.thinker
elif self.model_stage == "talker":
# Initialize talker model
talker_vllm_config = vllm_config.with_hf_config(
config.talker_config,
architectures=["Qwen3OmniMoeTalkerForConditionalGeneration"]
)
self.talker = init_vllm_registered_model(
vllm_config=talker_vllm_config,
prefix=maybe_prefix(prefix, "talker"),
hf_config=config.talker_config,
architectures=["Qwen3OmniMoeTalkerForConditionalGeneration"],
)
self.model = self.talker
elif self.model_stage == "code2wav":
# Initialize code2wav model
code2wav_vllm_config = vllm_config.with_hf_config(
config.code2wav_config,
architectures=["Qwen3OmniMoeCode2Wav"]
)
self.code2wav = init_vllm_registered_model(
vllm_config=code2wav_vllm_config,
prefix=maybe_prefix(prefix, "code2wav"),
hf_config=config.code2wav_config,
architectures=["Qwen3OmniMoeCode2Wav"],
)
self.model = self.code2wav
else:
raise ValueError(
f"Invalid model_stage: {self.model_stage}. "
f"Must be one of: 'thinker', 'talker', 'code2wav'"
)
```
#### Key Methods to Implement
1. **`forward()`**: Handles the forward pass for each stage
2. **`embed_input_ids()`**: Embeds input token IDs
3. **`embed_multimodal()`**: Processes multimodal inputs (if applicable)
4. **`compute_logits()`**: Computes logits from hidden states
5. **`load_weights()`**: Loads model weights with proper prefixing of different stages
### Step 4: Create `__init__.py`
Export the main model class:
```python
# vllm_omni/model_executor/models/qwen3_omni/__init__.py
from .qwen3_omni import Qwen3OmniMoeForConditionalGeneration
__all__ = ["Qwen3OmniMoeForConditionalGeneration"]
```
## Key Components
### 1. Model Interfaces
Your model should implement the appropriate interfaces:
- **`SupportsMultiModal`**: For models that process multimodal inputs
- **`SupportsPP`**: For models that support pipeline parallelism
- **`SupportsMRoPE`**: For models using multi-dimensional RoPE (if applicable)
### 2. Multimodal Registration
If your model processes multimodal inputs, register it with the multimodal registry:
```python
@MULTIMODAL_REGISTRY.register_processor(
YourMultiModalProcessor,
info=YourProcessingInfo,
dummy_inputs=YourDummyInputsBuilder,
)
class YourModel(nn.Module, SupportsMultiModal):
pass
```
### 3. Weight Loading
Implement `load_weights()` to handle weight loading with proper prefixing:
```python
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
"""Load weights for all components of the omni model."""
loaded_weights = set()
thinker_weights = []
talker_weights = []
code2wav_weights = []
# Separate weights by component
for k, v in weights:
if k.startswith("thinker."):
thinker_weights.append((k, v))
elif k.startswith("talker."):
talker_weights.append((k, v))
elif k.startswith("code2wav."):
code2wav_weights.append((k, v))
# Load each component's weights
if self.thinker and thinker_weights:
thinker_loaded = self.thinker.load_weights(thinker_weights)
thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, "thinker")
loaded_weights.update(thinker_loaded)
# Similar for talker and code2wav...
return loaded_weights
```
### 4. Output Format
Use `OmniOutput` for stage outputs:
```python
from vllm_omni.model_executor.models.output_templates import OmniOutput
# In forward method
return OmniOutput(
text_hidden_states=hidden_states,
multimodal_outputs={"additional_data": data},
next_token_id=next_token_id,
)
```
## Model Registration
Register your model in `vllm_omni/model_executor/models/registry.py`:
```python
_OMNI_MODELS = {
# ... existing models ...
# Your new model
"YourModelForConditionalGeneration": (
"your_model_name", # Module folder name
"your_model", # Module file name (without .py)
"YourModelForConditionalGeneration", # Class name
),
"YourModelThinkerForConditionalGeneration": (
"your_model_name",
"your_model_thinker",
"YourModelThinkerForConditionalGeneration",
),
# ... other stages ...
}
```
The registry uses lazy loading, so the model class is imported only when needed.
## Stage Configuration
Create a YAML configuration file in `vllm_omni/model_executor/stage_configs/`. For a complete example, see the [Qwen3-Omni configuration file](gh-file:vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml).
### Key Configuration Fields
- **`model_stage`**: Which stage to run ("thinker", "talker", "code2wav", etc.)
- **`model_arch`**: The model architecture name (must match registry)
- **`engine_input_source`**: List of stage IDs that provide input to this stage
- **`custom_process_input_func`**: Function to process inputs from previous stages
- **`final_output`**: Whether this stage produces the final output (True/False)
- **`final_output_type`**: Type of final output ("text", "audio", "image", etc.)
## Stage Input Processors
Stage transitions are the mechanism by which outputs from one stage are converted into inputs for the next stage. This section explains where and how stage transitions occur.
### Where Stage Transitions Are Called
Stage transitions happen automatically in the orchestrator (`OmniLLM` class) during the generation loop. Here's the detailed flow:
1. **Location**: `vllm_omni/entrypoints/omni_llm.py` in the `_run_generation()` method
2. **Trigger**: When a stage completes processing and produces outputs
3. **Execution Flow**:
```python
# In omni_llm.py, _run_generation() method (around line 345-460)
# Main orchestrator loop polls each stage for completed requests
for stage_id, stage in enumerate(self.stage_list):
result = stage.try_collect() # Get completed request
if result is None:
continue
# Store outputs from this stage
engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
stage.set_engine_outputs(engine_outputs)
# Check if there's a next stage to forward to
next_stage_id = stage_id + 1
if next_stage_id < len(self.stage_list):
next_stage: OmniStage = self.stage_list[next_stage_id]
# THIS IS WHERE STAGE TRANSITION HAPPENS
next_inputs = next_stage.process_engine_inputs(
self.stage_list,
[request_id_to_prompt[req_id]]
)
# Submit to next stage
task = {
"type": OmniStageTaskType.GENERATE,
"request_id": req_id,
"engine_inputs": next_inputs[0],
"sampling_params": sampling_params_list[next_stage_id],
}
next_stage.submit(task)
```
### How Stage Transitions Work
The stage transition process follows these steps:
1. **Stage Completion**: When a stage finishes processing a request, it stores outputs via `stage.set_engine_outputs(engine_outputs)`
2. **Transition Detection**: The orchestrator checks if there's a next stage and calls `process_engine_inputs()` on it
3. **Input Processing**: The `process_engine_inputs()` method in `OmniStage` (`omni_stage.py`) handles the transition:
```python
def process_engine_inputs(
self, stage_list: list[Any], prompt: OmniTokensPrompt | TextPrompt = None
) -> list[OmniTokensPrompt | TextPrompt]:
"""Process engine inputs for this stage from upstream stage outputs."""
if self.custom_process_input_func is None:
# Default behavior: pass token IDs directly
# Extract outputs from source stage
source_stage_id = self.engine_input_source[0]
source_outputs = stage_list[source_stage_id].engine_outputs
# ... create OmniTokensPrompt from token_ids ...
else:
# Custom transition function (YOUR CODE HERE)
return self.custom_process_input_func(
stage_list,
self.engine_input_source,
prompt,
self.requires_multimodal_data
)
```
- If `custom_process_input_func` is configured, it calls that function
- Otherwise, it uses default behavior (passing token IDs directly)
4. **Custom Function Execution**: Your custom function receives:
- `stage_list`: List of all stage objects (to access upstream stage outputs)
- `engine_input_source`: List of source stage IDs (e.g., `[0]` for stage 0)
- `prompt`: Original prompt data (for preserving multimodal data)
- `requires_multimodal_data`: Whether multimodal data is required
5. **Output Format**: The function must return a list of `OmniTokensPrompt` objects ready for the next stage
### Data Structures in Stage Transitions
Understanding the data structures is crucial for implementing stage transitions:
**Input to your function:**
- `stage_list[source_stage_id].engine_outputs`: List of `EngineCoreOutput` objects
- Each contains `outputs`: List of `RequestOutput` objects
- Each `RequestOutput` has:
- `token_ids`: Generated token IDs
- `multimodal_output`: Dict with keys like `"code_predictor_codes"`, etc.
- These are the hidden states or intermediate outputs from the model's forward pass
- `prompt_token_ids`: Original prompt token IDs
**Output from your function:**
- Must return `list[OmniTokensPrompt]` where each `OmniTokensPrompt` contains:
- `prompt_token_ids`: List[int] - Token IDs for the next stage
- `additional_information`: Dict[str, Any] - Optional metadata (e.g., embeddings, hidden states)
- `multi_modal_data`: Optional multimodal data if needed
### How Model Outputs Are Stored
The model's `forward()` method returns an `OmniOutput` object that contains:
- `text_hidden_states`: Final hidden states for text generation
- `multimodal_outputs`: Dict containing intermediate outputs
These outputs are captured during the forward pass and stored in `multimodal_output` with specific keys:
```python
# In your model's forward() method (e.g., qwen3_omni.py)
def forward(self, ...):
# ... processing ...
# For thinker stage: capture embeddings and hidden states
multimodal_outputs = {
"0": captured_embeddings, # Layer 0 embeddings
"24": captured_hidden_states, # Layer 24 hidden states
"tts_bos_embed": tts_bos_embed,
"tts_eos_embed": tts_eos_embed,
# ... other intermediate outputs ...
}
return OmniOutput(
text_hidden_states=hidden_states,
multimodal_outputs=multimodal_outputs,
)
```
These keys are then accessible in your stage transition function:
```python
# In stage_input_processors/qwen3_omni.py
thinker_embeddings = output.multimodal_output["0"] # Access by key
thinker_hidden_states = output.multimodal_output["24"]
```
### Key Points
1. **Accessing Upstream Outputs**: Use `stage_list[source_stage_id].engine_outputs` to get outputs from the source stage
2. **Extracting Data**: Access `output.multimodal_output[key]` to get specific hidden states or intermediate results
- Keys are defined by your model's `forward()` method when it creates `multimodal_outputs`
3. **Device Management**: Move tensors to appropriate devices (CPU for serialization, GPU for processing)
4. **Shape Transformations**: Reshape tensors as needed for the next stage (e.g., flattening codec codes)
5. **Batch Handling**: Process each request in the batch separately and return a list
### Complete Flow Diagram
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-dataflow-between-stages.png">
<img alt="Data Flow between stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-dataflow-between-stages.png" width=55%>
</picture>
</p>
### Implementation Example
Create stage transition processors in `vllm_omni/model_executor/stage_input_processors/your_model_name.py`:
```python
# qwen3_omni.py
def thinker2talker(
stage_list: list[Any],
engine_input_source: list[int],
prompt: OmniTokensPrompt | TextPrompt | None = None,
requires_multimodal_data: bool = False,
) -> list[OmniTokensPrompt]:
"""
Process thinker outputs to create talker inputs.
Args:
stage_list: List of stage objects
engine_input_source: Source stage IDs (typically [0] for thinker)
prompt: Original prompt data
Returns:
List of OmniTokensPrompt for talker stage
"""
source_stage_id = engine_input_source[0]
thinker_outputs = stage_list[source_stage_id].engine_outputs
talker_inputs = []
for thinker_output in thinker_outputs:
output = thinker_output.outputs[0]
# Extract thinker embeddings and hidden states
thinker_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda()
info = {
"thinker_embeddings": thinker_embeddings,
"thinker_hidden_states": thinker_hidden_states,
"thinker_sequences": thinker_output.prompt_token_ids + output.token_ids,
"thinker_input_ids": thinker_output.prompt_token_ids,
}
talker_inputs.append(
OmniTokensPrompt(
prompt_token_ids=[0] * computed_length,
additional_information=info,
multi_modal_data=None,
)
)
return talker_inputs
def talker2code2wav(
stage_list: list[Any],
engine_input_source: list[int],
prompt: OmniTokensPrompt | TextPrompt | None = None,
requires_multimodal_data: bool = False,
) -> list[OmniTokensPrompt]:
"""
Process talker outputs to create code2wav inputs.
"""
source_stage_id = engine_input_source[0]
talker_outputs = stage_list[source_stage_id].engine_outputs
code2wav_inputs = []
for talker_output in talker_outputs:
output = talker_output.outputs[0]
# Extract codec codes
codec_codes = (
output.multimodal_output["code_predictor_codes"]
.to(torch.long)
.transpose(0, 1)
.cpu()
.to(torch.long)
.reshape(-1)
.tolist()
)
code2wav_inputs.append(
OmniTokensPrompt(
prompt_token_ids=codec_codes,
multi_modal_data=None,
)
)
return code2wav_inputs
```
## Testing
For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
## Adding a Model Recipe
After implementing and testing your model, please add a model recipe to the [vllm-project/recipes](https://github.com/vllm-project/recipes) repository. This helps other users understand how to use your model with vLLM-Omni.
### What to Include
Your recipe should include:
1. **Model Overview**: Brief description of the model and its capabilities
2. **Installation Instructions**: Step-by-step setup instructions including:
- Installing vllm-omni and dependencies
- Installing any additional required packages (e.g., xformers, diffusers)
- Any version requirements
3. **Usage Examples**: Command-line examples demonstrating how to run the model
4. **Configuration Details**: Important configuration parameters and their meanings
### Example
For reference, see the [LongCat recipe example](https://github.com/vllm-project/recipes/pull/179) which demonstrates the expected format and structure.
### Recipe Location
Create your recipe file in the appropriate directory structure:
- For organization-specific models: `OrganizationName/ModelName.md`
- For general models: `ModelName.md`
The recipe should be a Markdown file that provides clear, reproducible instructions for users to get started with your model.
## Summary
Adding a new model to vLLM-Omni involves:
1. **Create model directory structure** with stage implementations
2. **Implement unified model class** that orchestrates stages
3. **Register model** in `registry.py`
4. **Create stage configuration** YAML file
5. **Implement stage input processors** for stage transitions
6. **Write tests** to verify functionality
7. **Add model recipe** to the [vllm-project/recipes](https://github.com/vllm-project/recipes) repository (see [Adding a Model Recipe](#adding-a-model-recipe) section)
### Qwen3-Omni Reference Files
For a complete reference implementation, see:
- **Main model**: `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py`
- **Thinker**: `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py`
- **Talker**: `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py`
- **Code2Wav**: `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py`
- **Stage config**: `vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml`
- **Input processors**: `vllm_omni/model_executor/stage_input_processors/qwen3_omni.py`
- **Registry**: `vllm_omni/model_executor/models/registry.py`
- **Testing**: `vllm_omni/tests/e2e/offline_inference/test_qwen3_omni.py`
For more information, see:
- [Architecture Overview](../../design/architecture_overview.md)
- [Supported Models](../../models/supported_models.md)
- [Stage Configuration Guide](../../configuration/stage_configs.md)
+166
View File
@@ -0,0 +1,166 @@
# Profiling vLLM-Omni
> **Warning:** Profiling incurs significant overhead. Use only for development and debugging, never in production.
vLLM-Omni uses the PyTorch Profiler to analyze performance across both **Multi-Stage LLMs** and **Diffusion Models**.
### 1. Set the Output Directory
Before running any script, set this environment variable. The system detects this and automatically saves traces here.
```bash
export VLLM_TORCH_PROFILER_DIR=./profiles
```
### 2. Start Profiling
It is best to limit profiling to one iteration to keep trace files manageable.
```bash
export VLLM_PROFILER_MAX_ITERS=1
```
**Selective Stage Profiling**
The profiler is default to function across all stages. But It is highly recommended to profile specific stages by passing the stages list, preventing from producing too large trace files:
```python
# Profile all stages
omni_llm.start_profile()
# Only profile Stage 1
omni_llm.start_profile(stages=[1])
```
```python
# Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
omni_llm.start_profile(stages=[0, 2])
```
**Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
```python
from vllm_omni import omni_llm
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
# 1. Start profiling if enabled
if profiler_enabled:
omni_llm.start_profile(stages=[0])
# Initialize generator
omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
total_requests = len(prompts)
processed_count = 0
# Main Processing Loop
for stage_outputs in omni_generator:
# ... [Output processing logic for text/audio would go here] ...
# Update count to track when to stop profiling
processed_count += len(stage_outputs.request_output)
# 2. Check if all requests are done to stop the profiler safely
if profiler_enabled and processed_count >= total_requests:
print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
# Stop the profiler while workers are still active
omni_llm.stop_profile()
# Wait for traces to flush to disk
print("[Info] Waiting 30s for workers to write trace files to disk...")
time.sleep(30)
print("[Info] Trace export wait time finished.")
omni_llm.close()
```
**Examples**:
1. **Qwen-omni 2.5**: [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
2. **Qwen-omni 3.0**: [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
**For Diffusion Models as a single stage**
Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding.
**CLI Usage:**
```python
python image_to_video.py \
--model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
--image qwen-bear.png \
--prompt "A cat playing with yarn, smooth motion" \
\
# Minimize Spatial Dimensions (Optional but helpful):
# Drastically reduces memory usage so the profiler doesn't
# crash due to overhead, though for accurate performance
# tuning you often want target resolutions.
--height 48 \
--width 64 \
\
# Minimize Temporal Dimension (Frames):
# Video models process 3D tensors (Time, Height, Width).
# Reducing frames to the absolute minimum (2) keeps the
# tensor size small, ensuring the trace file doesn't become
# multi-gigabytes in size.
--num_frames 2 \
\
# Minimize Iteration Loop (Steps):
# This is the most critical setting for profiling.
# Diffusion models run the same loop X times.
# Profiling 2 steps gives you the exact same performance
# data as 50 steps, but saves minutes of runtime and
# prevents the trace viewer from freezing.
--num_inference_steps 2 \
\
--guidance_scale 5.0 \
--guidance_scale_high 6.0 \
--boundary_ratio 0.875 \
--flow_shift 12.0 \
--fps 16 \
--output i2v_output.mp4
```
**Examples**:
1. **Qwen image edit**: [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**: [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
**Online Inference(Async)**
For online serving using AsyncOmni, the methods are asynchronous. This allows you to toggle profiling dynamically without restarting the server.
```python
from vllm_omni import AsyncOmni
# Inside an async function:
async_omni = AsyncOmni.from_engine_args(engine_args)
await async_omni.start_profile()
async for output in async_omni.generate(prompt, sampling_params, request_id):
# Process outputs...
pass
await async_omni.stop_profile()
```
### 3. Analyzing Omni Traces
Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
**Output**
**Chrome Trace** (```.json.gz```): Visual timeline of kernels and stages. Open in Perfetto UI.
**Viewing Tools:**
- [Perfetto](https://ui.perfetto.dev/)(recommended)
- ```chrome://tracing```(Chrome only)
**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM. See the official vLLM profiler documentation: [vLLM Profiling Guide](https://docs.vllm.ai/en/latest/dev/profiling.html)
+195
View File
@@ -0,0 +1,195 @@
# Architecture Overview
This document outlines the architectural design for vLLM-Omni.
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/omni-modality-model-architecture.png">
<img alt="Omni-Modality Model Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/omni-modality-model-architecture.png" width=55%>
</picture>
</p>
# Goals
The primary goal of the vLLM-Omni project is to build the fastest and easiest-to-use open-source Omni-Modality model inference & serving engine. vLLM-Omni extends the original vLLM, which was created to support large language models for text-based autoregressive (AR) generation tasks. vLLM-Omni is designed to support:
* **Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, and video, alongside text.
* **Non-Autoregressive Structure:** Support model structure beyond autoregressive, especially Diffusion Transformer (DiT), which is widely used in visual and audio generation.
* **Integration with vLLM Core:** Maintain compatibility and leverage existing vLLM key modules and optimizations where applicable.
* **Extensibility:** Design a modular and flexible architecture that can easily accommodate new modalities, model architectures, and output formats.
# Representative omni-modality models
According to analysis for current popular open-source models, most of them have the combination of AR+DiT. Specifically, they can be further categorized into 3 types below:
**DiT as a main structure, with AR as text encoder (e.g.: Qwen-Image)**
A powerful image generation foundation model capable of complex text rendering and precise image editing.
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/dit-main-architecture.png">
<img alt="Qwen-Image" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/dit-main-architecture.png" width=30%>
</picture>
</p>
**AR as a main structure, with DiT as multi-modal generator (e.g. BAGEL)**
A unified multimodal comprehension and generation model, with cot text output and visual generation.
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-main-architecture.png">
<img alt="Bagel" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-main-architecture.png" width=30%>
</picture>
</p>
**AR+DiT (e.g. Qwen-Omni)**
A natively end-to-end omni-modal LLM for multimodal inputs (text/image/audio/video...) and outputs (text/audio...).
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-dit-main-architecture.png">
<img alt="Qwen-Omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-dit-main-architecture.png" width=30%>
</picture>
</p>
# vLLM-Omni main architecture
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-main-architecture.png">
<img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-main-architecture.png" width=55%>
</picture>
</p>
## Key Components
| Component | Description |
| ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
| **OmniRouter** | provide an intelligent router for Omni-modality requests dispatch |
| **EntryPoints** | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni) and provide the OmniStage abstraction for different AR/DiT stages |
| **AR** | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management |
| **Diffusion** | natively implemented and optimized using acceleration components |
| **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages |
Disaggregated stages are managed through configuration, such as in the Qwen3-Omni example, where stages like Thinker, Talker, and Code2wav are defined as separate OmniStage instances with specific resources and input/output type.
## Main features
vLLM-Omni aims to be fast, flexible, and easy to use with the following features:
### Performance and Acceleration
The framework achieves high performance through several optimization techniques:
* **Efficient AR Support:** Leverages efficient KV cache management inherited from vLLM.
* **Pipelined Execution:** Uses pipelined stage execution overlapping to ensure high throughput.
* **Full Disaggregation:** Relies on the OmniConnector and dynamic resource allocation across stages.
* **Diffusion Acceleration:** Includes integrated support for diffusion acceleration. This is managed by the acceleration layer, which handles:
* **Cache:** Includes DBCache, TeaCache and third-party integration(e.g., [cache-dit](https://github.com/vipshop/cache-dit)).
* **Parallelism:** Supports TP, CP, USP, and CFG.
* **Attention:** Provides an interface for third-party integration (e.g., FA3, SAGE, MindIE-SD).
* **Quantization:** Supports various quantization implementations including FP8 and AWQ.
* **FusedOps:** Allows for custom and third-party integration.
### Flexibility and Usability
vLLM-Omni is designed to be flexible and straightforward for users:
* **Heterogeneous Pipeline Abstraction:** Manages complex model workflows effectively.
* **Hugging Face Integration:** Offers seamless integration with popular Hugging Face models.
* **Distributed Inference:** Supports tensor, pipeline, data, and expert parallelism.
* **Streaming Outputs:** Supports streaming outputs.
* **Unified API:** Provides a consistent and unified API interface compatible with vLLM.
* **OpenAI-compatible API Server:** Includes a FastAPI-based server for online serving that is compatible with the OpenAI API.
# Interface design
If you use vLLM, then you know how to use vLLM-Omni from Day 0:
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-user-interface.png">
<img alt="vLLM-Omni interface design" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-user-interface.png" width=55%>
</picture>
</p>
Taking **Qwen3-Omni** as an example:
## Offline Inference
The **Omni** class provides a Python interface for offline batched inference. Users initialize the Omni class with a Hugging Face model name and use the generate method, passing inputs that include both text prompts and multi-modal data:
```
# Create an omni_lm with HF model name.
from vllm_omni.entrypoints.omni import Omni
omni_lm = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
# Example prompts.
om_inputs = {"prompt": prompt,
"multi_modal_data": {
"video": video_frames,
"audio": audio_signal,
}}
# Generate texts and audio from the multi-modality inputs.
outputs = omni_lm.generate(om_inputs, sampling_params_list)
```
## Online Serving
Similar to vLLM, vLLM-Omni also provides a FastAPI-based server for online serving. Users can launch the server using the vllm serve command with the `--omni` flag:
```
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
Users can send requests to the server using curl:
```
# prepare user content
user_content='[
{
"type": "video_url",
"video_url": {
"url": "'"$SAMPLE_VIDEO_URL"'"
}
},
{
"type": "text",
"text": "Why is this video funny?"
}
]'
sampling_params_list='[
'"$thinker_sampling_params"',
'"$talker_sampling_params"',
'"$code2wav_sampling_params"'
]'
mm_processor_kwargs="{}"
# send the request
curl -sS -X POST http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d @- <<EOF
{
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"sampling_params_list": $sampling_params_list,
"mm_processor_kwargs": $mm_processor_kwargs,
"messages": [
{
"role": "system",
"content": [
{
"type": "text",
"text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
}
]
},
{
"role": "user",
"content": $user_content
}
]
}
```
For more usages, please refer to [examples](https://github.com/vllm-project/vllm-omni/tree/main/examples).
@@ -0,0 +1,106 @@
# Disaggregated Inference for Omni-Modality Models
This guide explains how to configure and use distributed connectors
(`vllm_omni/distributed/omni_connectors`) in vllm-omni for multi-stage pipelines.
Backend-specific setup lives in separate docs:
- [SharedMemoryConnector](omni_connectors/shared_memory_connector.md)
- [MooncakeConnector](omni_connectors/mooncake_connector.md)
- [YuanrongConnector](omni_connectors/yuanrong_connector.md)
## Overview
Connectors enable data transfer between pipeline stages (e.g., Thinker -> Talker).
Current connectors operate in D2H2D (device to host to device) mode.
## Connector Choices
| Use Case | Recommended Connector | Notes |
| :--- | :--- | :--- |
| Single node | SharedMemoryConnector | Auto-configured if no connector is specified. |
| Multi node (Mooncake) | MooncakeConnector | Requires Mooncake Master + metadata server. |
| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
## Core API
The connector system is built around `OmniConnectorBase`.
```python
class OmniConnectorBase(ABC):
@abstractmethod
def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, Optional[dict]]:
"""
Store data.
Returns: (success, serialized_size, metadata)
"""
pass
@abstractmethod
def get(self, from_stage: str, to_stage: str, get_key: str, metadata: Optional[dict] = None) -> Optional[tuple[Any, int]]:
"""
Retrieve data.
Args: metadata - transport-specific handles returned by put() (e.g., SHM name).
Returns: (object, serialized_size)
"""
pass
```
### Metadata Passing
Some connectors (e.g., SharedMemoryConnector) generate transient resources during `put()`.
This `metadata` must be passed through the control plane so `get()` can locate the data.
## Configuration Model
Define connectors in runtime:
```yaml
runtime:
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
```
Wire stages to connectors:
```yaml
stage_args:
- stage_id: 0
output_connectors:
to_stage_1: connector_of_shared_memory
- stage_id: 1
input_connectors:
from_stage_0: connector_of_shared_memory
```
If a pipeline edge has no explicit connector, the system auto-creates a
SharedMemoryConnector for that edge.
## Relationship with vLLM
vLLM provides specialized distributed mechanisms for specific artifacts:
- KV Transfer (`vllm.distributed.kv_transfer`): optimized for KV caches.
- EC Transfer (`vllm.distributed.ec_transfer`): optimized for encoder embeddings.
- Device Communicators (`vllm.distributed.device_communicators`): low-level primitives (NCCL, SHM).
vllm-omni complements this with a generalized connector abstraction:
1. Unifies transport via a single `put`/`get` API for any stage artifact.
2. Enables DAG-style pipelines across processes or nodes with per-edge transports.
3. Can wrap vLLM-specific transfers for KV paths while keeping a consistent interface.
## Operational Notes
- Fail-fast config validation: missing expected edges cause startup failures.
- Missing payloads halt stages: verify connector wiring and metadata propagation.
## Future Roadmap: D2D Transport
Current connectors use D2H2D paths. Future versions will introduce direct
device-to-device connectors (NCCL, UCX, IPC) to reduce latency for large
tensor payloads.
@@ -0,0 +1,74 @@
# MooncakeConnector
## When to Use
Best for multi-node distributed inference using Mooncake.
## Installation
```bash
# For CUDA-enabled systems (recommended)
pip install mooncake-transfer-engine
# For non-CUDA systems
pip install mooncake-transfer-engine-non-cuda
```
## Start Mooncake Master
```bash
# If you use Mooncake SSD storage
mkdir -p ./mc_storage
mooncake_master \
--rpc_port=50051 \
--enable_http_metadata_server=true \
--http_metadata_server_host=0.0.0.0 \
--http_metadata_server_port=8080 \
--metrics_port=9003 \
--root_fs_dir=./mc_storage/ \
--cluster_id=mc-local-1 &
```
## Configuration
Define the connector in runtime:
```yaml
runtime:
connectors:
connector_of_mooncake:
name: MooncakeConnector
extra:
host: "127.0.0.1"
metadata_server: "http://<MASTER_IP>:8080/metadata"
master: "<MASTER_IP>:50051"
segment: 512000000
localbuf: 64000000
proto: "tcp"
```
Wire stages to the connector:
```yaml
stage_args:
- stage_id: 0
output_connectors:
to_stage_1: connector_of_mooncake
- stage_id: 1
input_connectors:
from_stage_0: connector_of_mooncake
```
Parameters:
- host: local worker IP registered in the metadata server.
- metadata_server: metadata server URL for discovery and setup.
- master: Mooncake Master address.
- segment: global memory segment size in bytes.
- localbuf: local buffer size in bytes.
- proto: transport protocol ("tcp" or "rdma").
For more details, refer to the
[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
@@ -0,0 +1,27 @@
# SharedMemoryConnector
## When to Use
Best for single-node deployments where stages run on the same host. It is
auto-configured when no explicit connector is specified for an edge.
## How It Works
- Small payloads (< threshold): serialized and passed inline in metadata.
- Large payloads (>= threshold): stored in shared memory; the SHM name is
returned in metadata.
## Configuration
```yaml
runtime:
connectors:
connector_of_shared_memory:
name: SharedMemoryConnector
extra:
shm_threshold_bytes: 65536
```
## Notes
- Auto-mode uses SharedMemoryConnector if no connector is declared for an edge.
@@ -0,0 +1,101 @@
# YuanrongConnector
## When to Use
Best for multi-node distributed inference using Yuanrong Datasystem.
## Mechanism
Uses Yuanrong Datasystem's distributed KV store (`datasystem.kv_client`).
- Data Plane: TCP or RDMA for high-bandwidth transfer.
- Control Plane: Yuanrong Datasystem workers and etcd.
- Keying: deterministic keys based on `put_key` (often composed as `request_id:fromStage_toStage`).
## Installation
```bash
pip install openyuanrong-datasystem
```
## Start etcd
```bash
# Download and install etcd (v3.5.12 or higher)
ETCD_VERSION="v3.5.12"
ETCD_ARCH="linux-arm64"
wget https://github.com/etcd-io/etcd/releases/download/${ETCD_VERSION}/etcd-${ETCD_VERSION}-${ETCD_ARCH}.tar.gz
tar -xvf etcd-${ETCD_VERSION}-${ETCD_ARCH}.tar.gz
cd etcd-${ETCD_VERSION}-${ETCD_ARCH}
sudo cp etcd etcdctl /usr/local/bin/
# Start etcd
etcd \
--name etcd-single \
--data-dir /tmp/etcd-data \
--listen-client-urls http://0.0.0.0:2379 \
--advertise-client-urls http://0.0.0.0:2379 \
--listen-peer-urls http://0.0.0.0:2380 \
--initial-advertise-peer-urls http://0.0.0.0:2380 \
--initial-cluster etcd-single=http://0.0.0.0:2380 &
# Verify etcd is running
etcdctl --endpoints "127.0.0.1:2379" put key "value"
etcdctl --endpoints "127.0.0.1:2379" get key
```
For production environments, refer to the
[official etcd clustering documentation](https://etcd.io/docs/current/op-guide/clustering/).
## Start Datasystem Worker
```bash
# Replace ${ETCD_IP} with etcd node IP, ${WORKER_IP} with local node IP
dscli start -w \
--worker_address "${WORKER_IP}:31501" \
--etcd_address "${ETCD_IP}:2379" \
--shared_memory_size_mb 20480
```
To stop the worker:
```bash
dscli stop --worker_address "${WORKER_IP}:31501"
```
## Configuration
Define the connector in runtime:
```yaml
runtime:
connectors:
connector_of_yuanrong:
name: YuanrongConnector
extra:
host: "127.0.0.1"
port: 31501
get_sub_timeout_ms: 1000
```
Wire stages to the connector:
```yaml
stage_args:
- stage_id: 0
output_connectors:
to_stage_1: connector_of_yuanrong
- stage_id: 1
input_connectors:
from_stage_0: connector_of_yuanrong
```
Parameters:
- host: datasystem worker host.
- port: datasystem worker port.
- get_sub_timeout_ms: get timeout in milliseconds (0 for no timeout).
For more details, refer to the
[Yuanrong Datasystem repository](https://atomgit.com/openeuler/yuanrong-datasystem).
@@ -0,0 +1,59 @@
# Distributed utils
This directory (vllm_omni/distributed/ray_utils) contains utilities for distributed execution in vllm-omni, supporting both **Ray** and **Multiprocessing** backends.
## 1. Installation
```bash
pip install "ray[default]"
```
## 2. Ray Utils
The `ray_utils` module provides helper functions for managing Ray clusters and actors, which is essential for:
* **Multi-node deployment**: Running pipeline stages across different physical machines.
* **Resource management**: Efficient GPU/CPU allocation.
### 2.1 Basic Usage
To use the Ray backend, specify `worker_backend="ray"` when initializing the engine.
**Command Line Example:**
```bash
vllm serve Qwen/Qwen2.5-Omni-7B \
--omni \
--port 8091 \
--worker-backend ray \
--ray-address auto
```
### 2.2 Cluster Setup
**Step 1: Start Head Node**
Run this on your primary machine:
```bash
ray start --head --port=6399
```
**Step 2: Connect Worker Nodes**
Run this on each worker machine:
```bash
ray start --address=<HEAD_NODE_IP>:6399
```
> **Tip**: For a complete cluster setup script, refer to the vLLM example:
> [run_cluster.sh](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/run_cluster.sh)
### 2.3 Distributed Connector Support
When running on Ray, the system automatically adapts its communication strategy:
* **Cross-Node**: Recommended to use `MooncakeConnector` (requires separate configuration).
* **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
* **SHM threshold default differs**: when `worker_backend="ray"`, the SharedMemoryConnector default threshold is set to `sys.maxsize`, which forces payloads to go inline (no SHM). Override `shm_threshold_bytes` in the connector config if you want SHM for Ray runs.
### 2.4 Internal Helpers
* **`initialize_ray_cluster`**: Connects to an existing Ray cluster or starts a local one.
## 3. Troubleshooting
* **Connection Issues**: Ensure the Ray head node is accessible and ports (default 6399 in this example) are open.
* **Version Mismatch**: Ensure all nodes run the same version of Ray and Python.
+18
View File
@@ -0,0 +1,18 @@
# Design Documents
This section contains design documents and architecture specifications for vLLM-Omni.
## Architecture Documents
- [Architecture Overview](architecture_overview.md)
## Feature Design Documents
- [Disaggregated Inference](feature/disaggregated_inference.md)
- [Ray-based Execution](feature/ray_based_execution.md)
## Module Design Documents
- [AR Module](module/ar_module.md)
- [DIT Module](module/dit_module.md)
- [Entrypoint Module](module/entrypoint_module.md)
+414
View File
@@ -0,0 +1,414 @@
# AutoRegressive (AR) Module
## 1. Overview
The AutoRegressive (AR) module in vLLM-Omni handles autoregressive generation stages, primarily used for text, chain-of-thought(COT), and audio latent tokens generation stages in multi-stage models like Qwen2.5-Omni, Qwen3-Omni, BAGEL, .etc. Unlike some representative non-autoregressive generation stages (e.g., Diffusion), AR stages generate tokens sequentially, one at a time, following the standard transformer decoder pattern.
The AR module of vLLM-Omni extends vLLM's core components to support:
- **Multimodal inputs/outputs**: Processing images, videos, and audio alongside text
- **Direct embedding transfer**: Passing pre-computed prompt embeddings between pipeline stages via serialized payloads
- **Additional information flow**: Carrying per-request metadata (tensors, lists) through the pipeline
- **Hidden state exposure**: Exposing per-request hidden representations for downstream stages
- **Basic generator support**: Support some basic heterogeneous architecture such as Convolution, LSTM, etc.
As shown in the [end2end example](../../user_guide/examples/offline_inference/qwen3_omni.md), AR module can be widely applied across multiple stages, generating text tokens in thinker(AR), audio latent tokens in talker(AR) and audio wave in code2wav(Convolution).
## 2. Relationship with vLLM
The AR module builds upon vLLM main framework through inheritance, extending core classes while preserving compatibility with vLLM's scheduling, batching, KV cache management, and execution mechanisms.
### Inheritance Hierarchy
- Scheduler
```mermaid
classDiagram
class VLLMScheduler {
+schedule() SchedulerOutput
+update_from_output() EngineCoreOutputs
}
class OmniARScheduler {
+schedule() SchedulerOutput
}
class OmniGenerationScheduler {
+schedule() SchedulerOutput
+update_from_output() EngineCoreOutputs
}
VLLMScheduler <|-- OmniARScheduler
VLLMScheduler <|-- OmniGenerationScheduler
```
- Worker
```mermaid
classDiagram
class GPUWorker {
+init_device()
+model_runner
}
class GPUARWorker {
+init_device()
}
class GPUGenerationWorker {
+init_device()
}
GPUWorker <|-- GPUARWorker
GPUWorker <|-- GPUGenerationWorker
```
- ModelRunner
```mermaid
classDiagram
class GPUModelRunner {
+execute_model()
+sample_tokens()
}
class OmniGPUModelRunner {
+_update_states()
+_preprocess()
+_model_forward()
}
class GPUARModelRunner {
+execute_model()
+sample_tokens()
}
class GPUGenerationModelRunner {
+execute_model()
}
GPUModelRunner <|-- OmniGPUModelRunner
OmniGPUModelRunner <|-- GPUARModelRunner
OmniGPUModelRunner <|-- GPUGenerationModelRunner
```
- InputProcessor/OutputProcessor
```mermaid
classDiagram
class InputProcessor {
+process_inputs() EngineCoreRequest
}
class OmniInputProcessor {
+process_inputs() OmniEngineCoreRequest
}
class VLLMOutputProcessor {
+process_outputs() OutputProcessorOutput
}
class MultimodalOutputProcessor {
+process_outputs() OutputProcessorOutput
+_route_and_normalize()
}
InputProcessor <|-- OmniInputProcessor
VLLMOutputProcessor <|-- MultimodalOutputProcessor
```
### Key Extensions
- **Scheduler**: `OmniARScheduler` extends `vllm.v1.core.sched.scheduler.Scheduler` to enrich scheduled requests with omni-specific payloads
- **Worker**: `GPUARWorker` extends `vllm.v1.worker.gpu_worker.Worker` to initialize AR-specific model runners
- **ModelRunner**: `GPUARModelRunner` extends `OmniGPUModelRunner` → `vllm.v1.worker.gpu_model_runner.GPUModelRunner` to expose hidden states and handle multimodal outputs
- **InputProcessor**: `OmniInputProcessor` extends `vllm.v1.engine.input_processor.InputProcessor` to serialize prompt embeddings and additional information
- **OutputProcessor**: `MultimodalOutputProcessor` extends `vllm.v1.engine.output_processor.OutputProcessor` to route and accumulate multimodal outputs
## 3. Scheduler Design
The AR module provides two scheduler implementations: one for standard autoregressive generation and one for basic heterogeneous architectures.
### Request Flow
The following diagram illustrates the request flow through the AR module components:
```mermaid
flowchart TD
A[OmniInputProcessor] -->|OmniEngineCoreRequest| B[OmniARScheduler]
B -->|schedule: OmniNewRequestData| C[GPUARWorker]
C -->|SchedulerOutput| D[GPUARModelRunner]
D -->|execute_model: None| E[Model Forward Pass]
E -->|hidden_states, logits| D
D -->|sample_tokens: OmniModelRunnerOutput| F[OmniARScheduler]
F -->|update_from_output| G[MultimodalOutputProcessor]
G -->|RequestOutput| H[Client/Downstream Stage]
style A fill:#e1f5ff
style B fill:#fff4e1
style C fill:#e8f5e9
style D fill:#f3e5f5
style G fill:#fce4ec
```
The flow follows vLLM's standard pattern: input processing → scheduling → worker execution → output processing, with omni-specific enrichments at each stage.
### OmniARScheduler
`OmniARScheduler` extends the base vLLM scheduler with minimal modifications, focusing on enriching scheduled requests with omni-specific payloads.
#### Modified API: `schedule()`
The scheduler wraps base `NewRequestData` entries with `OmniNewRequestData` to include prompt embeddings and additional information:
```python
def schedule(self) -> SchedulerOutput:
scheduler_output = super().schedule()
# Rewrap base NewRequestData entries with OmniNewRequestData
new_list = []
for nr in scheduler_output.scheduled_new_reqs:
request = self.requests.get(nr.req_id)
omni_nr = OmniNewRequestData(
req_id=nr.req_id,
prompt_token_ids=nr.prompt_token_ids,
# ... other base fields ...
prompt_embeds=getattr(request, "prompt_embeds", None),
additional_information=getattr(request, "additional_information", None),
)
new_list.append(omni_nr)
scheduler_output.scheduled_new_reqs = new_list
return scheduler_output
```
The `update_from_output()` method remains unchanged, inheriting standard request lifecycle management from the base scheduler.
### OmniGenerationScheduler
`OmniGenerationScheduler` implements a fast-path scheduling strategy for basic heterogeneous architectures that process all input tokens in a single step.
#### Modified API: `schedule()`
Allocates all input tokens for a request at once (or 1 placeholder if zero), falling back to default scheduling if budget is insufficient:
```python
def schedule(self) -> SchedulerOutput:
# Fast path: allocate all input tokens at once
while self.waiting and token_budget > 0:
request = self.waiting.peek_request()
required_tokens = max(getattr(request, "num_prompt_tokens", 0), 1)
if required_tokens > token_budget:
break # Fall back to default scheduling
# Allocate and schedule...
```
#### Modified API: `update_from_output()`
Marks requests as finished immediately after one step, since generation models complete in a single forward pass:
```python
def update_from_output(self, ...) -> dict[int, EngineCoreOutputs]:
# ...
# Diffusion request: completes in one step
request.status = RequestStatus.FINISHED_STOPPED
kv_transfer_params = self._free_request(request)
# ...
```
## 4. Worker and ModelRunner Design
### GPUARWorker
`GPUARWorker` initializes the AR-specific model runner while maintaining standard device initialization:
```python
class GPUARWorker(GPUWorker):
def init_device(self):
# ... standard device initialization ...
self.model_runner = GPUARModelRunner(self.vllm_config, self.device)
```
### GPUARModelRunner
`GPUARModelRunner` follows vLLM's two-phase execute/sample flow while exposing hidden states and multimodal outputs.
#### Two-Phase Execution
**Phase 1: `execute_model()`** - Runs forward pass and stores state:
- Computes logits from hidden states
- Stores `ExecuteModelState` with hidden states, logits, and multimodal outputs
- Returns `None` to defer sampling
**Phase 2: `sample_tokens()`** - Samples tokens and builds output:
- Retrieves stored state from `execute_model()`
- Samples tokens using logits
- Extracts per-request hidden states and multimodal outputs
- Builds `OmniModelRunnerOutput` with `pooler_output` containing hidden states
```python
def sample_tokens(self, grammar_output) -> OmniModelRunnerOutput:
# Retrieve stored state
hidden_states, multimodal_outputs = self.execute_model_state
# Sample tokens
sampler_output = self._sample(logits, spec_decode_metadata)
# Extract per-request hidden states
pooler_output = []
for rid in req_ids:
hidden_slice = hidden_states_cpu[start:end]
payload = {"hidden": hidden_slice}
# Add multimodal outputs if present
pooler_output.append(payload)
return OmniModelRunnerOutput(
pooler_output=pooler_output,
# ... other fields ...
)
```
### GPUGenerationModelRunner
`GPUGenerationModelRunner` implements a simplified single-phase execution for basic heterogeneous architectures:
- No logits computation or token sampling
- Direct generation from forward pass in model implementation
- Returns outputs via `pooler_output` immediately after forward pass
### OmniGPUModelRunner
`OmniGPUModelRunner` provides shared functionality for both AR and Generation runners:
#### Prompt Embeddings Overlay
During prefill, overlays custom `prompt_embeds` from request state onto `inputs_embeds`:
```python
def _collect_additional_information_for_prefill(self, num_scheduled_tokens_np):
for req_index, req_id in enumerate(self.input_batch.req_ids):
req_state = self.requests[req_id]
pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
# Overlay prompt_embeds for prefill portion
if pe_cpu is not None:
src = pe_cpu[num_computed_tokens:num_computed_tokens + overlay_len]
self.inputs_embeds[start_offset:start_offset + overlay_len].copy_(src)
```
#### Additional Information Processing
Decodes and manages `additional_information` payloads:
- Decodes serialized payloads → CPU tensors in request state
- Passes runtime information to model via `runtime_additional_information` kwarg
- Processes model-provided updates via `postprocess()` hook
- Merges updates back into request state
#### M-RoPE Position Initialization
For multimodal models using M-RoPE (e.g., Qwen2-VL), computes position encodings from multimodal feature metadata (image grids, video grids, audio features).
## 5. Input/Output Processing
### Processing Pipeline
The input/output processing pipeline handles serialization, routing, and accumulation of multimodal data:
```mermaid
sequenceDiagram
participant Client
participant OmniInputProcessor
participant Scheduler
participant ModelRunner
participant MultimodalOutputProcessor
participant Client
Client->>OmniInputProcessor: prompt + prompt_embeds + additional_info
OmniInputProcessor->>OmniInputProcessor: Serialize tensors to payloads
OmniInputProcessor->>Scheduler: OmniEngineCoreRequest (with payloads)
Scheduler->>ModelRunner: OmniNewRequestData (with payloads)
ModelRunner->>ModelRunner: Decode payloads → CPU tensors
ModelRunner->>ModelRunner: Overlay prompt_embeds on inputs_embeds
ModelRunner->>ModelRunner: Forward pass with runtime_additional_information
ModelRunner->>ModelRunner: Extract hidden states + multimodal outputs
ModelRunner->>MultimodalOutputProcessor: OmniModelRunnerOutput (pooler_output)
MultimodalOutputProcessor->>MultimodalOutputProcessor: Route by output_type
MultimodalOutputProcessor->>MultimodalOutputProcessor: Accumulate tensors in OmniRequestState
MultimodalOutputProcessor->>MultimodalOutputProcessor: Consolidate tensor lists
MultimodalOutputProcessor->>Client: RequestOutput (with multimodal_output)
```
### OmniInputProcessor
`OmniInputProcessor` extends the base input processor to serialize prompt embeddings and additional information for inter-stage transfer.
#### Payload Serialization
Converts PyTorch tensors to serialized payloads:
```python
def process_inputs(self, ...) -> OmniEngineCoreRequest:
# Serialize prompt_embeds
if "prompt_embeds" in decoder_inputs:
pe_cpu = decoder_inputs["prompt_embeds"].detach().to("cpu").contiguous()
prompt_embeds_payload = PromptEmbedsPayload(
data=pe_cpu.numpy().tobytes(),
shape=[seq_len, hidden_size],
dtype=dtype_str,
)
# Serialize additional_information
if "additional_information" in decoder_inputs:
entries = {}
for key, value in raw_info.items():
if isinstance(value, torch.Tensor):
entry = AdditionalInformationEntry(
tensor_data=value.numpy().tobytes(),
tensor_shape=list(value.shape),
tensor_dtype=dtype_str,
)
entries[key] = entry
additional_information_payload = AdditionalInformationPayload(entries=entries)
return OmniEngineCoreRequest(
# ... standard fields ...
prompt_embeds=prompt_embeds_payload,
additional_information=additional_information_payload,
)
```
### MultimodalOutputProcessor
`MultimodalOutputProcessor` routes outputs by modality type and accumulates multimodal tensors.
#### Output Routing
Routes `EngineCoreOutput` by `output_type` attribute:
- `"text"`: Standard text generation path
- `"image"`, `"audio"`, `"latents"`: Extract from `pooling_output` or `multimodal_outputs`
- Fallback: Heuristic based on presence of `pooling_output`
#### Tensor Accumulation
`OmniRequestState` accumulates multimodal tensors across multiple steps:
```python
def add_multimodal_tensor(self, payload, mm_type):
# Normalize payload to dict
incoming = {mm_type or "hidden": payload}
# Accumulate: convert tensors to lists for deferred concatenation
if isinstance(v, torch.Tensor) and isinstance(existing, torch.Tensor):
self.mm_accumulated[k] = [existing, v] # List accumulation
```
Before final output, consolidates tensor lists via concatenation:
```python
def _consolidate_multimodal_tensors(self):
for k, v in self.mm_accumulated.items():
if isinstance(v, list) and isinstance(v[0], torch.Tensor):
self.mm_accumulated[k] = torch.cat(v, dim=0) # Concatenate
```
The consolidated tensors are attached to `RequestOutput.multimodal_output` for consumption by downstream stages or clients.
## 6. Summary
The AR module of vLLM-Omni extends vLLM through strategic inheritance and minimal API modifications:
### Key Design Patterns
1. **Inheritance over composition**: Extends vLLM classes to preserve compatibility with existing scheduling, batching, and execution mechanisms
2. **Payload serialization**: Uses `PromptEmbedsPayload` and `AdditionalInformationPayload` for efficient inter-stage data transfer
3. **Two-phase execution**: Maintains vLLM's execute/sample separation for AR models while supporting single-phase execution for generation models
4. **Multimodal routing**: Routes outputs by `output_type` and accumulates tensors incrementally to support streaming
### Differences from vLLM
- **Payload support**: Serialized prompt embeddings and additional information enable direct transfer between pipeline stages
- **Multimodal handling**: Extended input/output processors support images, audio, and other modalities alongside text
- **Hidden state exposure**: AR model runners expose per-request hidden states via `pooler_output` for downstream consumption
- **Generation scheduler**: Fast-path scheduling for basic heterogeneous architectures that complete in one step
The AR module seamlessly integrates with vLLM's existing infrastructure while adding the necessary extensions for multi-stage, multimodal generation pipelines.
+906
View File
@@ -0,0 +1,906 @@
---
toc_depth: 4
---
# Diffusion Module Architecture Design
The vLLM-Omni diffusion module (`vllm_omni/diffusion`) is a high-performance inference engine for diffusion models, designed with a modular architecture that separates concerns across multiple components. It provides efficient execution for non-autoregressive generation tasks such as image and video generation.
This document describes the architecture design of the diffusion module, including the diffusion engine, scheduler, worker, diffusion pipeline, and acceleration components.
<p align="center">
<img src="https://github.com/user-attachments/assets/cde1ebea-4006-4d30-8283-15a5721a09d2" alt="vLLM-Omni Diffusion Module Components" width="80%">
</p>
<p align="center">
<em> Main Components of the Diffusion Module </em>
</p>
**Table of Content:**
- [Architecture Overview](#architecture-overview)
- [Diffusion Engine](#1-diffusion-engine)
- [Scheduler](#2-scheduler)
- [Worker](#3-worker)
- [Diffusion Pipeline](#4-diffusion-pipeline)
- [Acceleration Components](#5-acceleration-components)
- [Attention Backends](#51-attention-backends)
- [Parallel Attention](#52-parallel-attention)
- [Cache Backends](#53-cache-backends)
- [Parallel Strategies](#54-parallel-strategies)
- [Data Flow](#6-data-flow)
---
## Architecture Overview
The diffusion module follows a **multi-process, distributed architecture** with clear separation of concerns:
<p align="center">
<img src="https://github.com/user-attachments/assets/78c4d446-d238-4406-a057-eb17d2ccc8e0" alt="vLLM-Omni Diffusion Module Architecture" width="100%">
</p>
<p align="center">
<em> Diffusion Architecture Overview </em>
</p>
---
## 1. Diffusion Engine
**Location**: `vllm_omni/diffusion/diffusion_engine.py`
### Responsibilities
The `DiffusionEngine` is the **orchestrator** of the diffusion inference system. It manages the lifecycle of worker processes and coordinates the execution flow.
### Key Components
#### 1.1 Initialization
```python
class DiffusionEngine:
def __init__(self, od_config: OmniDiffusionConfig):
self.od_config = od_config
self.post_process_func = get_diffusion_post_process_func(od_config)
self.pre_process_func = get_diffusion_pre_process_func(od_config)
self._processes: list[mp.Process] = []
self._make_client()
```
**Key Features**:
- **Pre/Post Processing**: Registers model-specific pre-processing and post-processing functions via registry pattern
- **Worker Management**: Launches and manages multiple worker processes (one per GPU)
- **Process Isolation**: Uses multiprocessing for true parallelism
#### 1.2 Worker Launch Process
The engine launches workers using a **spawn** method:
```python
def _launch_workers(self, broadcast_handle):
# Creates one process per GPU
for i in range(num_gpus):
process = mp.Process(
target=worker_proc.worker_main,
args=(i, od_config, writer, broadcast_handle),
name=f"DiffusionWorker-{i}",
)
process.start()
```
**Design Decisions**:
- **Spawn Method**: Ensures clean state for each worker (no shared memory issues)
- **Pipe Communication**: Uses `mp.Pipe` for initialization handshake
- **Device Selection**: Each worker is assigned a specific GPU (`cuda:{rank}`)
#### 1.3 Request Processing Flow
```python
def step(self, requests: list[OmniDiffusionRequest]):
# 1. Pre-process requests
requests = self.pre_process_func(requests)
# 2. Send to scheduler and wait for response
output = self.add_req_and_wait_for_response(requests)
# 3. Post-process results
result = self.post_process_func(output.output)
return result
```
**Flow**:
1. **Pre-processing**: Applies model-specific transformations
2. **Scheduling**: Delegates to scheduler for distribution
3. **Post-processing**: Converts raw outputs to final format (e.g., PIL images)
---
## 2. Scheduler
**Location**: `vllm_omni/diffusion/scheduler.py`
### Architecture
The `Scheduler` is implemented as a **Singleton** pattern to ensure a single coordination point across the system, i.e., only one scheduler instance exists for coordination.
### Key Components
#### 2.1 Message Queue System
```python
class Scheduler:
def initialize(self, od_config: OmniDiffusionConfig):
# Broadcast queue: scheduler -> all workers
self.mq = MessageQueue(
n_reader=od_config.num_gpus,
n_local_reader=od_config.num_gpus,
local_reader_ranks=list(range(od_config.num_gpus)),
)
# Result queue: rank 0 worker -> scheduler
self.result_mq = None # Initialized later
```
**Communication Pattern**:
- **Broadcast Queue**: One-to-many communication (scheduler → all workers)
- **Result Queue**: One-to-one communication (rank 0 → scheduler)
- **Shared Memory**: Uses `MessageQueue` (ZMQ-based) for efficient IPC
#### 2.2 Request Distribution
```python
def add_req(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
# Broadcast request to all workers
self.mq.enqueue(requests)
# Wait for result from Rank 0
output = self.result_mq.dequeue()
return output
```
**Design Features**:
- **Broadcast Model**: All workers receive the same request (for tensor parallelism)
- **Single Response**: Only rank 0 sends results back (avoids duplicate outputs)
- **Synchronous**: Blocks until result is received (can be made async)
#### 2.3 Singleton Pattern
```python
class Scheduler:
_instance = None
def __new__(cls, *args, **kwargs):
if not cls._instance:
cls._instance = super().__new__(cls)
return cls._instance
# Global singleton instance
scheduler = Scheduler()
```
**Benefits**:
- **Single Point of Control**: Ensures consistent state
- **Easy Access**: Global `scheduler` instance accessible everywhere
- **Resource Management**: Centralized queue management
---
## 3. Worker
**Location**: `vllm_omni/diffusion/worker/gpu_worker.py`
### Architecture
Workers are **independent processes** that execute the actual model inference. Each worker runs on a dedicated GPU and participates in distributed inference.
### Key Components
#### 3.1 Worker Process Structure
```python
class WorkerProc:
def __init__(self, od_config, gpu_id, broadcast_handle):
# Initialize ZMQ context for IPC
self.context = zmq.Context(io_threads=2)
# Connect to broadcast queue (receive requests)
self.mq = MessageQueue.create_from_handle(broadcast_handle, gpu_id)
# Create result queue (only rank 0)
if gpu_id == 0:
self.result_mq = MessageQueue(n_reader=1, ...)
# Initialize GPU worker
self.worker = GPUWorker(local_rank=gpu_id, rank=gpu_id, od_config=od_config)
```
**Initialization Steps**:
1. **IPC Setup**: Creates ZMQ context and message queues
2. **Distributed Environment Setup**: Initializes PyTorch distributed communication
- For CUDA GPUs: Uses NCCL (fast GPU communication)
- For NPU: Uses HCCL (Huawei Collective Communications Library)
- For other devices: Uses appropriate backend (GLOO, MCCL, etc.)
3. **Model Loading**: Loads diffusion pipeline on assigned GPU
4. **Cache Setup**: Enables cache backend if configured.
#### 3.2 GPU Worker
```python
class GPUWorker:
def init_device_and_model(self):
# Set distributed environment variables
os.environ["RANK"] = str(rank)
os.environ["WORLD_SIZE"] = str(world_size)
# Initialize PyTorch distributed
init_distributed_environment(world_size, rank)
parallel_config = self.od_config.parallel_config
initialize_model_parallel(
data_parallel_size=parallel_config.data_parallel_size,
cfg_parallel_size=parallel_config.cfg_parallel_size,
sequence_parallel_size=parallel_config.sequence_parallel_size,
tensor_parallel_size=parallel_config.tensor_parallel_size,
pipeline_parallel_size=parallel_config.pipeline_parallel_size,
)
# Load model
model_loader = DiffusersPipelineLoader(load_config)
self.pipeline = model_loader.load_model(od_config, load_device=f"cuda:{rank}")
# Setup cache backend
from vllm_omni.diffusion.cache.selector import get_cache_backend
self.cache_backend = get_cache_backend(od_config.cache_backend, od_config.cache_config)
if self.cache_backend is not None:
self.cache_backend.enable(self.pipeline)
```
**Key Features**:
- **Tensor Parallelism**: Supports multi-GPU tensor parallelism via PyTorch distributed
- **Model Loading**: Uses `DiffusersPipelineLoader` for efficient weight loading
- **Cache Integration**: Enables cache backends (TeaCache, cache-dit, etc.) transparently
#### 3.3 Worker Busy Loop
```python
def worker_busy_loop(self):
while self._running:
# 1. Receive unified message (generation request, RPC request, or shutdown)
msg = self.recv_message()
# 2. Route message based on type
if isinstance(msg, dict) and msg.get("type") == "rpc":
# Handle RPC request
result, should_reply = self.execute_rpc(msg)
if should_reply:
self.return_result(result)
elif isinstance(msg, dict) and msg.get("type") == "shutdown":
# Handle shutdown message
self._running = False
else:
# Handle generation request (OmniDiffusionRequest list)
output = self.worker.execute_model(msg, self.od_config)
self.return_result(output)
```
**Execution Flow**:
1. **Receive**: Dequeues unified messages from shared memory queue
2. **Route**: Handles different message types (generation, RPC, shutdown)
3. **Execute**: Runs forward pass through pipeline for generation requests
4. **Respond**: Sends results back (rank 0 for generation, specified rank for RPC)
#### 3.4 Model Execution
```python
@torch.inference_mode()
def execute_model(self, reqs: list[OmniDiffusionRequest], od_config):
req = reqs[0] # TODO: support batching
# Refresh cache backend if enabled
if self.cache_backend is not None and self.cache_backend.is_enabled():
self.cache_backend.refresh(self.pipeline, req.num_inference_steps)
# Set forward context for parallelism
with set_forward_context(
vllm_config=self.vllm_config,
omni_diffusion_config=self.od_config
):
output = self.pipeline.forward(req)
return output
```
The model execution leverages multiple parallelism strategies that are transparently applied during the forward pass. The `set_forward_context()` context manager makes parallel group information available throughout the forward pass:
```python
# Inside transformer layers, parallel groups are accessed via:
from vllm_omni.diffusion.distributed.parallel_state import (
get_sp_group, get_dp_group, get_cfg_group, get_pp_group
)
```
**Optimizations**:
- **Cache Refresh**: Clears cache state before each generation for clean state
- **Context Management**: Forward context ensures parallel groups are available during execution
- **Single Request**: Currently processes one request at a time (batching TODO)
---
## 4. Diffusion Pipeline
**Location**: `vllm_omni/diffusion/models/*/pipeline_*.py`
The pipeline is the **model-specific implementation** that orchestrates the diffusion process. Different models (QwenImage, Wan2.2, Z-Image) have their own pipeline implementations.
Most pipeline implementation are referred from `diffusers`. The multi-step diffusion loop is usually the most time-consuming part during the overall inference process, which is defined by the `diffuse` function in the pipeline class. An example is as follows:
```python
def diffuse(self, ...):
for i, t in enumerate(timesteps):
# Forward pass for positive prompt
transformer_kwargs = {
"hidden_states": latents,
"timestep": timestep / 1000,
"encoder_hidden_states": prompt_embeds,
}
noise_pred = self.transformer(**transformer_kwargs)[0]
# Forward pass for negative prompt (CFG)
if do_true_cfg:
neg_transformer_kwargs = {...}
neg_transformer_kwargs["cache_branch"] = "negative"
neg_noise_pred = self.transformer(**neg_transformer_kwargs)[0]
# Combine predictions
comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
noise_pred = comb_pred * (cond_norm / noise_norm)
# Scheduler step
latents = self.scheduler.step(noise_pred, t, latents)[0]
return latents
```
**Key Features**:
- **CFG Support**: Handles classifier-free guidance with separate forward passes
- **Cache Branching**: Uses `cache_branch` parameter for cache-aware execution
- **True CFG**: Implements advanced CFG with norm preservation
To learn more about the diffusion pipeline and how to add a new diffusion pipeline, please view [Adding Diffusion Model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_diffusion_model)
---
## 5. Acceleration Components
### 5.1 Attention Backends
**Location**: `vllm_omni/diffusion/attention/`
#### Architecture
The attention system uses a **backend selector pattern** that automatically chooses the optimal attention implementation based on hardware and model configuration.
#### Backend Selection
**Location**: `vllm_omni/diffusion/attention/selector.py`
```python
class Attention(nn.Module):
def __init__(self, num_heads, head_size, causal, softmax_scale, ...):
# Auto-select backend
self.attn_backend = get_attn_backend(-1)
self.attn_impl_cls = self.attn_backend.get_impl_cls()
self.attention = self.attn_impl_cls(...)
```
**Available Backends**:
- **FlashAttention**: Optimized CUDA kernel (FA2/FA3) - memory efficient via tiling
- **SDPA**: PyTorch's scaled dot-product attention - default, cross-platform
- **SageAttention**: Sparse attention implementation from SageAttention library
- **AscendAttention**: NPU-optimized attention for Ascend hardware
These backends provide the **kernel implementations** for attention computation. For attention-level sequence parallelism strategies (Ring Attention, Ulysses), see [Parallel Attention](#52-parallel-attention).
#### Backend Selection Mechanism
```python
def get_attn_backend(head_size: int) -> type[AttentionBackend]:
# Check environment variable
backend_name = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
if backend_name:
return load_backend(backend_name.upper())
# Default to SDPA
return SDPABackend
```
**Selection Priority**:
1. **Environment Variable**: `DIFFUSION_ATTENTION_BACKEND` for manual override
- Valid values: `FLASH_ATTN`, `TORCH_SDPA`, `SAGE_ATTN`, `ASCEND`
- Example: `export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN`
2. **Automatic Fallback**: Falls back to SDPA if selected backend unavailable
3. **Hardware Detection**: Can select based on device type (NPU, CUDA, etc.)
**Backend Availability**:
- **SDPA**: Always available (PyTorch built-in)
- **FlashAttention**: Requires `flash-attn` package installed
- **SageAttention**: Requires `sage-attention` package (from THU-ML GitHub)
- **AscendAttention**: Only available on Ascend NPU hardware
#### Attention Backend Registry
**Location**: `vllm_omni/diffusion/attention/selector.py`
The attention system uses a **registry pattern** to manage and dynamically load attention backends. This allows for easy extension and runtime selection of backends.
**Registry Structure**:
```python
# Registry mapping backend names to their module paths and class names
_BACKEND_CONFIG = {
"FLASH_ATTN": {
"module": "vllm_omni.diffusion.attention.backends.flash_attn",
"class": "FlashAttentionBackend",
},
"TORCH_SDPA": {
"module": "vllm_omni.diffusion.attention.backends.sdpa",
"class": "SDPABackend",
},
"SAGE_ATTN": {
"module": "vllm_omni.diffusion.attention.backends.sage_attn",
"class": "SageAttentionBackend",
},
"ASCEND": {
"module": "vllm_omni.diffusion.attention.backends.ascend_attn",
"class": "AscendAttentionBackend",
},
}
```
#### Attention Backend Integration
The `Attention` layer integrates backends through a unified interface. Here's how **FlashAttentionBackend** is integrated as an example:
```python
# attention/backends/flash_attn.py
class FlashAttentionBackend(AttentionBackend):
@staticmethod
def get_name() -> str:
return "FLASH_ATTN"
@staticmethod
def get_impl_cls() -> type["FlashAttentionImpl"]:
return FlashAttentionImpl
@staticmethod
def get_supported_head_sizes() -> list[int]:
return [64, 96, 128, 192, 256] # FlashAttention supports these head sizes
class FlashAttentionImpl(AttentionImpl):
def __init__(self, num_heads, head_size, softmax_scale, causal, ...):
self.num_heads = num_heads
self.causal = causal
self.softmax_scale = softmax_scale
def forward(self, query, key, value, attn_metadata=None):
# Call FlashAttention kernel
out = flash_attn_func(
query, key, value,
causal=self.causal,
softmax_scale=self.softmax_scale,
)
return out
```
---
### 5.2 Parallel Attention
**Location**: `vllm_omni/diffusion/attention/parallel/`
#### Architecture
Parallel attention strategies implement **Sequence Parallelism (SP) at the attention layer level**. These strategies distribute attention computation across multiple GPUs by splitting the sequence dimension, using different communication patterns. They work **on top of** AttentionBackend implementations (FlashAttention, SDPA, etc.), handling the parallelization/communication while the backends handle the actual attention computation.
**Key Distinction**: Unlike AttentionBackend (which provides kernel implementations), ParallelAttentionStrategy provides communication patterns for multi-GPU attention parallelism. These strategies implement the `ParallelAttentionStrategy` interface and use AttentionBackend implementations internally.
Both Ring Attention and Ulysses are forms of Sequence Parallelism (SP) that:
- Split the sequence dimension across GPUs
- Contribute to `sequence_parallel_size` (via `ring_degree` and `ulysses_degree`)
- Work at the attention layer level (not model/pipeline level)
#### Ulysses Sequence Parallelism (USP)
**Location**: `vllm_omni/diffusion/attention/parallel/ulysses.py`
USP is a sequence-parallel attention strategy that splits attention computation across multiple GPUs by distributing both the sequence dimension and attention heads. It uses **all-to-all communication** to efficiently parallelize attention for very long sequences. Specifically, it uses **all-to-all** collective operations to redistribute Q/K/V tensors before attention computation and gather results afterward.
Ulysses splits attention computation in two dimensions:
1. **Sequence Dimension**: Splits the sequence length across GPUs
2. **Head Dimension**: Splits attention heads across GPUs
**Configuration**: `ulysses_degree` contributes to `sequence_parallel_size`
#### Ring Sequence Parallelism
**Location**: `vllm_omni/diffusion/attention/parallel/ring.py`
Ring Attention is a **parallel attention strategy** that implements sequence parallelism using ring-based point-to-point (P2P) communication. Unlike attention backends that provide the attention kernel implementation, Ring Attention is a **communication pattern** that works on top of attention backends (FlashAttention or SDPA).
Ring Attention splits sequence dimension across GPUs in a ring topology, implemented via the `ParallelAttentionStrategy` interface, instead of `AttentionBackend`. P2P ring communication is applied to circulate Key/Value blocks across GPUs. Internally, `ring_flash_attn_func` or `ring_pytorch_attn_func` is used depending on available backends.
**Architecture**:
```python
class RingParallelAttention:
"""Ring sequence-parallel strategy."""
def run_attention(self, query, key, value, attn_metadata, ...):
# Selects underlying attention kernel (FlashAttention or SDPA)
if backend_pref == "sdpa":
return ring_pytorch_attn_func(...) # Uses SDPA kernel
else:
return ring_flash_attn_func(...) # Uses FlashAttention kernel
```
**Integration**:
- Ring Attention is activated when `ring_degree > 1` in parallel config
- It's selected by `build_parallel_attention_strategy()` in the attention layer
- The `Attention` layer routes to `_run_ring_attention()` when Ring is enabled
- Works alongside attention backends: Ring handles communication, backends handle computation
**Configuration**: `ring_degree` contributes to `sequence_parallel_size`
#### Relationship with AttentionBackend
Parallel attention strategies (Ring, Ulysses) work **on top of** AttentionBackend implementations:
- They use AttentionBackend for the actual attention computation (FlashAttention, SDPA, etc.)
- They handle the multi-GPU communication/parallelization layer
- They implement `ParallelAttentionStrategy` interface (not `AttentionBackend`)
For general parallelism strategies (Data Parallelism, Tensor Parallelism, Pipeline Parallelism), see [Parallel Strategies](#54-parallel-strategies).
---
### 5.3 Cache Backends
**Location**: `vllm_omni/diffusion/cache/`
#### Architecture
Cache backends provide a **unified interface** for applying different caching strategies to accelerate diffusion inference. The system supports multiple backends (TeaCache, cache-dit) with a consistent API for enabling and refreshing cache state.
#### Cache Backend Interface
```python
class CacheBackend(ABC):
def __init__(self, config: DiffusionCacheConfig):
self.config = config
self.enabled = False
@abstractmethod
def enable(self, pipeline: Any) -> None:
"""Enable cache on the pipeline."""
raise NotImplementedError
@abstractmethod
def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
"""Refresh cache state for new generation."""
raise NotImplementedError
def is_enabled(self) -> bool:
"""Check if cache is enabled."""
return self.enabled
```
**Design Pattern**:
- **Abstract Base Class**: Defines contract for all cache backends
- **Pipeline-based**: Works with pipeline instances (not just transformers)
- **State Management**: Provides refresh mechanism for clean state between generations
#### Available Backends
**1. TeaCache Backend**
**Location**: `vllm_omni/diffusion/cache/teacache/backend.py`
```python
class TeaCacheBackend(CacheBackend):
def enable(self, pipeline: Any):
# Extract transformer from pipeline
transformer = pipeline.transformer
transformer_type = transformer.__class__.__name__
# Create TeaCacheConfig from DiffusionCacheConfig
teacache_config = TeaCacheConfig(
transformer_type=transformer_type,
rel_l1_thresh=self.config.rel_l1_thresh,
coefficients=self.config.coefficients,
)
# Apply hooks to transformer
apply_teacache_hook(transformer, teacache_config)
self.enabled = True
def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True):
transformer = pipeline.transformer
if hasattr(transformer, "_hook_registry"):
transformer._hook_registry.reset_hook(TeaCacheHook._HOOK_NAME)
```
**TeaCache Features**:
- **Timestep-aware**: Caches based on timestep embedding similarity
- **Adaptive**: Dynamically decides when to reuse cached computations
- **CFG-aware**: Handles positive/negative branches separately
- **Custom Hook System**: Uses a custom forward interception mechanism (via `HookRegistry`) that wraps the module's `forward` method, allowing transparent integration without modifying model code
**2. Cache-DiT Backend**
**Location**: `vllm_omni/diffusion/cache/cache_dit_backend.py`
```python
class CacheDiTBackend(CacheBackend):
def enable(self, pipeline: Any):
# Uses cache-dit library for acceleration
# Supports DBCache, SCM (Step Computation Masking), TaylorSeer
# Works with single and dual-transformer architectures
...
self.enabled = True
def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True):
# Updates cache context with new num_inference_steps
...
```
**Cache-DiT Features**:
- **DBCache**: Dynamic block caching with configurable compute blocks
- **SCM**: Step Computation Masking for additional speedup
- **TaylorSeer**: Advanced calibration for cache accuracy
- **Dual-transformer Support**: Handles models like Wan2.2 with two transformers
#### Cache Backend Selector
**Location**: `vllm_omni/diffusion/cache/selector.py`
```python
def get_cache_backend(
cache_backend: str | None,
cache_config: dict | DiffusionCacheConfig
) -> CacheBackend | None:
"""Get cache backend instance based on cache_backend string.
Args:
cache_backend: Cache backend name ("cache_dit", "tea_cache", or None)
cache_config: Cache configuration (dict or DiffusionCacheConfig)
Returns:
Cache backend instance or None if cache_backend is "none"
"""
if cache_backend is None or cache_backend == "none":
return None
if isinstance(cache_config, dict):
cache_config = DiffusionCacheConfig.from_dict(cache_config)
if cache_backend == "cache_dit":
return CacheDiTBackend(cache_config)
elif cache_backend == "tea_cache":
return TeaCacheBackend(cache_config)
else:
raise ValueError(f"Unsupported cache backend: {cache_backend}")
```
**Usage Flow**:
1. **Selection**: `get_cache_backend()` returns appropriate backend instance
2. **Enable**: `backend.enable(pipeline)` called during worker initialization
3. **Refresh**: `backend.refresh(pipeline, num_inference_steps)` called before each generation
4. **Check**: `backend.is_enabled()` verifies cache is active
### 5.4 Parallel Strategies
**Location**: `vllm_omni/diffusion/distributed/parallel_state.py`
#### Parallelism Types
The system supports multiple orthogonal parallelism strategies:
**Sequence Parallelism (SP)**
- **Purpose**: Split sequence dimension across GPUs
- **Attention-level SP**: Ring Attention and Ulysses (USP) implement SP at the attention layer level
- See [Parallel Attention](#52-parallel-attention) for details
- Configuration: `ulysses_degree` × `ring_degree` = `sequence_parallel_size`
- **Use Case**: Very long sequences (e.g., high-resolution images)
**Data Parallelism (DP)**
- **Purpose**: Replicate model across GPUs, split batch
- **Use Case**: Batch processing, throughput optimization
**Tensor Parallelism (TP)** (Experimental)
- **Purpose**: Split model weights across GPUs
- **Implementation**: Uses vLLM's tensor parallel groups
- **Use Case**: Large models that don't fit on single GPU
**CFG Parallelism** (under development)
- **Purpose**: Parallelize Classifier-Free Guidance (positive/negative prompts)
- **Infrastructure**: CFG parallel groups are initialized and available via `get_cfg_group()`
#### Parallel Group Management
```python
def initialize_model_parallel(
data_parallel_size: int = 1,
cfg_parallel_size: int = 1,
sequence_parallel_size: int | None = None,
ulysses_degree: int = 1,
ring_degree: int = 1,
tensor_parallel_size: int = 1,
pipeline_parallel_size: int = 1,
vae_parallel_size: int = 0,
):
# Generate orthogonal parallel groups
rank_generator = RankGenerator(
tensor_parallel_size,
sequence_parallel_size,
pipeline_parallel_size,
cfg_parallel_size,
data_parallel_size,
"tp-sp-pp-cfg-dp",
)
# Initialize each parallel group
_DP = init_model_parallel_group(rank_generator.get_ranks("dp"), ...)
_CFG = init_model_parallel_group(rank_generator.get_ranks("cfg"), ...)
_SP = init_model_parallel_group(rank_generator.get_ranks("sp"), ...)
_PP = init_model_parallel_group(rank_generator.get_ranks("pp"), ...)
_TP = init_model_parallel_group(rank_generator.get_ranks("tp"), ...)
```
**Rank Order**: `tp-sp-pp-cfg-dp` (tensor → sequence → pipeline → cfg → data)
**Note**: For attention-level Sequence Parallelism implementations (Ring Attention and Ulysses), see [Parallel Attention](#52-parallel-attention). This section covers higher-level parallelism strategies.
---
## 6. Data Flow
### Complete Request Flow
<p align="center">
<img src="https://github.com/user-attachments/assets/6e093a0d-29c0-4efd-85d5-747002bd2fed" alt="vLLM-Omni Diffusion Module Components" width="100%">
</p>
<p align="center">
<em> End-to-end Data Flow in the vLLM-Omni Diffusion Module </em>
</p>
```
1. User Request
└─> OmniDiffusion.generate(prompt)
└─> Prepare OmniDiffusionRequest
└─> DiffusionEngine.step(requests)
2. Pre-processing
└─> pre_process_func(requests)
└─> Model-specific transformations
3. Scheduling
└─> scheduler.add_req(requests)
└─> Broadcast via MessageQueue to all workers
4. Worker Execution
└─> WorkerProc.worker_busy_loop()
└─> GPUWorker.execute_model(reqs)
└─> Pipeline.forward(req)
├─> encode_prompt()
├─> prepare_latents()
├─> diffuse() [loop]
│ ├─> transformer.forward() [with cache backend hooks]
│ └─> scheduler.step()
└─> vae.decode()
5. Result Collection
└─> Rank 0 sends DiffusionOutput via result queue
└─> Scheduler receives and returns
6. Post-processing
└─> post_process_func(output)
└─> Convert to PIL images / final format
```
---
+1
View File
@@ -0,0 +1 @@
Architecture design of the entrypoint (update soon)
+6
View File
@@ -0,0 +1,6 @@
# Examples
vLLM-Omni's examples are split into two categories:
- If you are using vLLM-Omni from within Python code, see the *Offline Inference* section.
- If you are using vLLM-Omni from an HTTP application or client, see the *Online Serving* section.
+39
View File
@@ -0,0 +1,39 @@
# Sleep Mode
vLLM-Omni’s **Sleep Mode** allows you to temporarily release most GPU memory used by a model—such as model weights and key-value (KV) caches (for autoregressive models)—**without stopping the server or unloading the Docker container**.
This feature is inherited from [vLLM’s Sleep Mode](https://blog.vllm.ai/2025/10/26/sleep-mode.html), which provides zero-reload model switching for multi-model serving.
It is especially useful in **RLHF**, **training**, or **cost-saving scenarios**, where GPU resources must be freed between inference workloads.
---
## Omni Model
Omni model inherit the feature from vLLM' Sleep Mode
This means:
- Support both Level 1 and Level 2 sleep, allow to release and reset both model weights and KV Cache
## Diffusion Model Extension
We added Sleep Mode support for **diffusion models**, which previously lacked this functionality.
In diffusion pipelines, this currently only offloads **model weight memory**, as these models typically do not use KV caches.
This means:
- Diffusion models can now enter Level 1 sleep.
- Pipeline states (e.g., noise schedulers, buffers) remain intact after waking.
- Useful for releasing VRAM between image generation or training cycles.
---
## Enable sleep mode
To enable sleep mode, set the `enable_sleep_mode` in `engine_args` to `True`
Example:
```python
omni = Omni(model=...,enable_sleep_mode=True)
```
@@ -0,0 +1,4 @@
nav:
- README.md
- gpu.md
- npu.md
@@ -0,0 +1,8 @@
# Installation
vLLM-Omni supports the following hardware platforms:
- [GPU](gpu.md)
- [NVIDIA CUDA](gpu.md)
- [AMD ROCm](gpu.md)
- [NPU](npu.md)
+66
View File
@@ -0,0 +1,66 @@
# GPU
vLLM-Omni is a Python library that supports the following GPU variants. The library itself mainly contains python implementations for framework and models.
## Requirements
- OS: Linux
- Python: 3.12
!!! note
vLLM-Omni is currently not natively supported on Windows.
=== "NVIDIA CUDA"
--8<-- "docs/getting_started/installation/gpu/cuda.inc.md:requirements"
=== "AMD ROCm"
--8<-- "docs/getting_started/installation/gpu/rocm.inc.md:requirements"
## Set up using Python
### Create a new Python environment
--8<-- "docs/getting_started/installation/python_env_setup.inc.md"
### Pre-built wheels
=== "NVIDIA CUDA"
--8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-wheels"
=== "AMD ROCm"
--8<-- "docs/getting_started/installation/gpu/rocm.inc.md:pre-built-wheels"
[](){ #build-from-source }
### Build wheel from source
=== "NVIDIA CUDA"
--8<-- "docs/getting_started/installation/gpu/cuda.inc.md:build-wheel-from-source"
=== "AMD ROCm"
--8<-- "docs/getting_started/installation/gpu/rocm.inc.md:build-wheel-from-source"
## Set up using Docker
### Pre-built images
=== "NVIDIA CUDA"
--8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-images"
=== "AMD ROCm"
--8<-- "docs/getting_started/installation/gpu/rocm.inc.md:pre-built-images"
### Build your own docker image
=== "AMD ROCm"
--8<-- "docs/getting_started/installation/gpu/rocm.inc.md:build-docker"
@@ -0,0 +1,103 @@
# --8<-- [start:requirements]
- GPU: compute capability 7.0 or higher (e.g., V100, T4, RTX20xx, A100, L4, H100, etc.)
# --8<-- [end:requirements]
# --8<-- [start:set-up-using-python]
vLLM-Omni depends vLLM. So please follow instructions below mainly for vLLM.
!!! note
PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See <gh-issue:8420> for more details.
In order to be performant, vLLM has to compile many cuda kernels. The compilation unfortunately introduces binary incompatibility with other CUDA versions and PyTorch versions, even for the same PyTorch version with different building configurations.
Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new** environment. If either you have a different CUDA version or you want to use an existing PyTorch installation, you need to build vLLM from source. See [build-from-source-vllm](https://docs.vllm.ai/en/stable/getting_started/installation/gpu/#build-wheel-from-source) for more details.
# --8<-- [start:pre-built-wheels]
#### Installation of vLLM
Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
vLLM-Omni is built based on vLLM. Please install it with command below.
```bash
uv pip install vllm==0.14.0 --torch-backend=auto
```
#### Installation of vLLM-Omni
```bash
uv pip install vllm-omni
```
# --8<-- [end:pre-built-wheels]
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.14.0 release version of the library
```bash
uv pip install vllm==0.14.0 --torch-backend=auto
```
The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
```bash
git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
uv pip install -e .
```
<details><summary>(Optional) Installation of vLLM from source</summary>
If you want to check, modify or debug with source code of vLLM, install the library from source with the following instructions:
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.14.0
```
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
```bash
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.14.0/vllm-0.14.0-cp38-abi3-manylinux_2_31_x86_64.whl
```
Install vllm with command below (If you have no existing PyTorch).
```bash
uv pip install --editable .
```
Install vllm with command below (If you already have PyTorch).
```bash
python use_existing_torch.py
uv pip install -r requirements/build.txt
uv pip install --no-build-isolation --editable .
```
</details>
# --8<-- [end:build-wheel-from-source]
# --8<-- [start:build-wheel-from-source-in-docker]
# --8<-- [end:build-wheel-from-source-in-docker]
# --8<-- [start:pre-built-images]
vLLM-Omni offers an official docker image for deployment. These images are built on top of vLLM docker images and available on Docker Hub as [vllm/vllm-omni](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which release of vLLM it is based on.
Here's an example deployment command that has been verified on 2 x H100's:
```bash
docker run --runtime nvidia --gpus 2 \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
--ipc=host \
vllm/vllm-omni:v0.14.0rc1 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
```
!!! tip
You can use this docker image to serve models the same way you would with in vLLM! To do so, make sure you overwrite the default entrypoint (`vllm serve --omni`) which works only for models supported in the vLLM-Omni project.
# --8<-- [end:pre-built-images]
@@ -0,0 +1,105 @@
# --8<-- [start:requirements]
- GPU: Validated on gfx942 (It should be supported on the AMD GPUs that are supported by vLLM.)
# --8<-- [end:requirements]
# --8<-- [start:set-up-using-python]
vLLM-Omni current recommends the steps in under setup through Docker Images.
# --8<-- [start:pre-built-wheels]
# --8<-- [end:pre-built-wheels]
# --8<-- [start:build-wheel-from-source]
# --8<-- [end:build-wheel-from-source]
# --8<-- [start:build-docker]
#### Build docker image
```bash
DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.rocm -t vllm-omni-rocm .
```
#### Launch the docker image
##### Launch with OpenAI API Server
```
docker run --rm \
--group-add=video \
--ipc=host \
--cap-add=SYS_PTRACE \
--security-opt seccomp=unconfined \
--device /dev/kfd \
--device /dev/dri \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
--ipc=host \
vllm-omni-rocm \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
```
##### Launch with interactive session for development
```
docker run --rm -it \
--network=host \
--group-add=video \
--ipc=host \
--cap-add=SYS_PTRACE \
--security-opt seccomp=unconfined \
--device /dev/kfd \
--device /dev/dri \
-v <path/to/model>:/app/model \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--entrypoint bash \
vllm-omni-rocm
```
# --8<-- [end:build-docker]
# --8<-- [start:pre-built-images]
vLLM-Omni offers an official docker image for deployment. These images are built on top of vLLM docker images and available on Docker Hub as [vllm/vllm-omni-rocm](https://hub.docker.com/r/vllm/vllm-omni-rocm/tags). The version of vLLM-Omni indicates which release of vLLM it is based on.
#### Launch vLLM-Omni Server
Here's an example deployment command that has been verified on 2 x MI300's:
```bash
docker run --rm \
--group-add=video \
--ipc=host \
--cap-add=SYS_PTRACE \
--security-opt seccomp=unconfined \
--device /dev/kfd \
--device /dev/dri \
-v <path/to/model>:/app/model \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
vllm/vllm-omni-rocm:v0.14.0rc1 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
#### Launch an interactive terminal with prebuilt docker image.
If you want to run in dev environment you can launch the docker image as follows:
```bash
docker run --rm -it \
--network=host \
--group-add=video \
--ipc=host \
--cap-add=SYS_PTRACE \
--security-opt seccomp=unconfined \
--device /dev/kfd \
--device /dev/dri \
-v <path/to/model>:/app/model \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--entrypoint bash \
vllm/vllm-omni-rocm:v0.14.0rc1
```
# --8<-- [end:pre-built-images]
+23
View File
@@ -0,0 +1,23 @@
# NPU
vLLM-Omni supports NPU through the vLLM Ascend Plugin (vllm-ascend). This is a community maintained hardware plugin for running vLLM on NPU.
## Requirements
- OS: Linux
- Python: 3.12
!!! note
vLLM-Omni is currently not natively supported on Windows.
=== "NPU"
--8<-- "docs/getting_started/installation/npu/npu.inc.md:requirements"
## Installation
### Recommended
=== "NPU"
--8<-- "docs/getting_started/installation/npu/npu.inc.md:installation"
@@ -0,0 +1,59 @@
# --8<-- [start:requirements]
For detailed hardware and software requirements, please refer to the [vllm-ascend installation documentation](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
# --8<-- [end:requirements]
# --8<-- [start:installation]
The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built Docker images:
```bash
# Update DEVICE according to your NPUs (/dev/davinci[0-7])
export DEVICE0=/dev/davinci0
export DEVICE1=/dev/davinci1
# Update the vllm-ascend image
# Atlas A2:
# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0
# Atlas A3:
# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0-a3
export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0
docker run --rm \
--name vllm-omni-npu \
--shm-size=1g \
--device $DEVICE0 \
--device $DEVICE1 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
-v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v /root/.cache:/root/.cache \
-p 8000:8000 \
-it $IMAGE bash
# Install the missing dependency of mooncake in the origin image.
apt update
apt install libjemalloc2
echo "export LD_PRELOAD=/usr/lib/$(uname -m)-linux-gnu/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc
source ~/.bashrc
# Inside the container, install vLLM-Omni from source
cd /vllm-workspace
git clone -b v0.14.0 https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
pip install -v -e .
export VLLM_WORKER_MULTIPROC_METHOD=spawn
# (Optional) Disable mooncake for stable capability
mv /usr/local/Ascend/ascend-toolkit/latest/python/site-packages/mooncake \
/usr/local/Ascend/ascend-toolkit/latest/python/site-packages/mooncake.disabled
```
The default workdir is `/workspace`, with vLLM, vLLM-Ascend and vLLM-Omni code placed in `/vllm-workspace` installed in development mode.
For other installation methods (pip installation, building from source, custom Docker builds), please refer to the [vllm-ascend installation guide](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
# --8<-- [end:installation]
@@ -0,0 +1,6 @@
It's recommended to use [uv](https://docs.astral.sh/uv/), a very fast Python environment manager, to create and manage Python environments. Please follow the [documentation](https://docs.astral.sh/uv/#getting-started) to install `uv`. After installing `uv`, you can create a new Python environment using the following commands:
```bash
uv venv --python 3.12 --seed
source .venv/bin/activate
```
+116
View File
@@ -0,0 +1,116 @@
# Quickstart
This guide will help you quickly get started with vLLM-Omni to perform:
- Offline batched inference
- Online serving using OpenAI-compatible server
## Prerequisites
- OS: Linux
- Python: 3.12
## Installation
For installation on GPU from source:
```bash
uv venv --python 3.12 --seed
source .venv/bin/activate
# On CUDA
uv pip install vllm==0.14.0 --torch-backend=auto
# On ROCm
uv pip install vllm==0.14.0 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
uv pip install -e .
```
For additional installation methods — please see the [installation guide](installation/README.md).
## Offline Inference
Text-to-image generation quickstart with vLLM-Omni:
```python
from vllm_omni.entrypoints.omni import Omni
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[0].images
images[0].save("coffee.png")
```
You can pass a list of prompts and wait for them to process altogether, shown below.
!!! info
However, it is not currently recommended to do so
because not all models support batch inference,
and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
```python
from vllm_omni.entrypoints.omni import Omni
if __name__ == "__main__":
omni = Omni(
model="Tongyi-MAI/Z-Image-Turbo",
# stage_configs_path="./stage-config.yaml", # See below
)
prompts = [
"a cup of coffee on a table",
"a toy dinosaur on a sandy beach",
"a fox waking up in bed and yawning",
]
omni_outputs = omni.generate(prompts)
for i_prompt, prompt_output in enumerate(omni_outputs):
this_request_output = prompt_output.request_output[0]
this_images = this_request_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")
# saved to p0-img0.jpg
# saved to p1-img0.jpg
# saved to p2-img0.jpg
```
!!! info
For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
batched inputs, you can [modify this configuration](../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
## Online Serving with OpenAI-Completions API
Text-to-image generation quickstart with vLLM-Omni:
```bash
vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091
```
```bash
curl -s http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"messages": [
{"role": "user", "content": "a cup of coffee on the table"}
],
"extra_body": {
"height": 1024,
"width": 1024,
"num_inference_steps": 50,
"guidance_scale": 4.0,
"seed": 42
}
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > coffee.png
```
For more details, please refer to [online serving](../user_guide/examples/online_serving/text_to_image.md).
+277
View File
@@ -0,0 +1,277 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Hook to automatically generate docs/api/README.md from the codebase.
This script scans the vllm_omni module for public classes and functions,
categorizes them, and generates a summary README file.
"""
import ast
import logging
from pathlib import Path
logger = logging.getLogger("mkdocs")
ROOT_DIR = Path(__file__).parent.parent.parent.parent
API_README_PATH = ROOT_DIR / "docs" / "api" / "README.md"
# Category mappings: module prefix -> category name and description
CATEGORIES = {
"entrypoints": {
"name": "Entry Points",
"description": "Main entry points for vLLM-Omni inference and serving.",
},
"inputs": {
"name": "Inputs",
"description": "Input data structures for multi-modal inputs.",
},
"engine": {
"name": "Engine",
"description": "Engine classes for offline and online inference.",
},
"core": {
"name": "Core",
"description": "Core scheduling and caching components.",
},
# "model_executor": {
# "name": "Model Executor",
# "description": "Model execution components.",
# },
"config": {
"name": "Configuration",
"description": "Configuration classes.",
},
"worker": {
"name": "Workers",
"description": "Worker classes and model runners for distributed inference.",
},
}
class APIVisitor(ast.NodeVisitor):
"""AST visitor to extract public classes and module-level functions."""
def __init__(self, module_path: str):
self.module_path = module_path
self.classes: list[str] = []
self.functions: list[str] = []
self._class_stack: list[str] = [] # Track nested class definitions
def visit_ClassDef(self, node: ast.ClassDef) -> None:
"""Visit class definitions."""
if not node.name.startswith("_"):
self.classes.append(f"{self.module_path}.{node.name}")
# Track that we're entering a class
self._class_stack.append(node.name)
self.generic_visit(node)
# Remove from stack when done visiting
self._class_stack.pop()
def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
"""Visit function definitions - only collect module-level functions."""
# Only collect if we're not inside a class (stack is empty)
if not self._class_stack and not node.name.startswith("_"):
self.functions.append(f"{self.module_path}.{node.name}")
self.generic_visit(node)
def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
"""Visit async function definitions - only collect module-level functions."""
# Only collect if we're not inside a class (stack is empty)
if not self._class_stack and not node.name.startswith("_"):
self.functions.append(f"{self.module_path}.{node.name}")
self.generic_visit(node)
def parse_file_for_symbols(file_path: Path, module_path: str) -> tuple[list[str], list[str]]:
"""
Parse a Python file and extract public classes and functions.
Returns:
Tuple of (classes, functions)
"""
try:
# If this is __init__.py, use parent module path
if file_path.name == "__init__.py":
# Remove .__init__ from module path
if module_path.endswith(".__init__"):
module_path = module_path[:-9]
with open(file_path, encoding="utf-8") as f:
content = f.read()
tree = ast.parse(content, filename=str(file_path))
visitor = APIVisitor(module_path)
visitor.visit(tree)
return visitor.classes, visitor.functions
except Exception as e:
logger.debug(f"Could not parse {file_path}: {e}")
return [], []
def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]:
"""
Scan the vllm_omni package and categorize public symbols.
Returns:
Dict mapping category names to lists of symbol full names
"""
categorized: dict[str, list[str]] = {cat["name"]: [] for cat in CATEGORIES.values()}
try:
# Find the package directory
package_path = ROOT_DIR / package_name
if not package_path.exists():
logger.warning(f"Package path not found: {package_path}")
return categorized
# Walk through all Python files
for py_file in package_path.rglob("*.py"):
# Skip __init__.py and private modules
if py_file.name.startswith("_") and py_file.name != "__init__.py":
continue
# Get module path
relative_path = py_file.relative_to(ROOT_DIR)
module_path = str(relative_path.with_suffix("")).replace("/", ".").replace("\\", ".")
# Skip excluded modules (avoid importing vllm during docs build)
excluded_prefixes = [
"vllm_omni.diffusion.models.qwen_image",
"vllm_omni.entrypoints.async_diffusion",
"vllm_omni.entrypoints.openai",
]
if any(module_path.startswith(prefix) for prefix in excluded_prefixes):
continue
# Handle __init__.py - use parent module path
if py_file.name == "__init__.py":
# Remove .__init__ from module path
if module_path.endswith(".__init__"):
module_path = module_path[:-9]
# Determine category from module path
category = None
for prefix, cat_info in CATEGORIES.items():
if prefix in module_path:
category = cat_info["name"]
break
if not category:
continue
# Parse file for symbols
classes, functions = parse_file_for_symbols(py_file, module_path)
# Filter out internal implementation classes
# Skip classes that look like internal components (DiT layers, etc.)
internal_patterns = [
"Block",
"Layer",
"Net",
"Embedding",
"Norm",
"Activation",
"Solver",
"Pooling",
"Attention",
"MLP",
"DecoderLayer",
"InputEmbedding",
"TimestepEmbedding",
"CodecEmbedding",
"DownSample",
"UpSample",
"Res2Net",
"SqueezeExcitation",
"TimeDelay",
"TorchActivation",
"SnakeBeta",
"SinusPosition",
"RungeKutta",
"AMPBlock",
"AdaLayerNorm",
]
# Add classes (filter out internal ones)
for class_name in classes:
class_short_name = class_name.split(".")[-1]
# Skip if it matches internal patterns (unless it's a main model class)
if any(pattern in class_short_name for pattern in internal_patterns):
# But include main model classes
if not any(
main_class in class_short_name
for main_class in [
"ForConditionalGeneration",
"Model",
"Registry",
"Worker",
"Runner",
"Scheduler",
"Manager",
"Processor",
"Config",
]
):
continue
categorized[category].append(class_name)
# Add important functions (parse, preprocess, etc.)
for func_name in functions:
# Include functions that match certain patterns
if any(keyword in func_name.lower() for keyword in ["parse", "preprocess"]):
categorized[category].append(func_name)
# Sort symbols within each category
for category in categorized:
categorized[category].sort()
except Exception as e:
logger.error(f"Error scanning package: {e}", exc_info=True)
return categorized
def generate_readme(categorized: dict[str, list[str]]) -> str:
"""Generate the API README markdown content."""
lines = ["# Summary", ""]
# Generate sections for each category
for prefix, cat_info in CATEGORIES.items():
category_name = cat_info["name"]
description = cat_info["description"]
symbols = categorized.get(category_name, [])
if not symbols:
continue
lines.append(f"## {category_name}")
lines.append("")
lines.append(description)
lines.append("")
for symbol in symbols:
lines.append(f"- [{symbol}][]")
lines.append("")
return "\n".join(lines)
def on_startup(command, dirty: bool):
"""MkDocs hook entry point."""
logger.info("Generating API README documentation")
# Scan the package
categorized = scan_package()
# Generate README content
content = generate_readme(categorized)
# Write to file
API_README_PATH.parent.mkdir(parents=True, exist_ok=True)
with open(API_README_PATH, "w", encoding="utf-8") as f:
f.write(content)
logger.info(f"API README generated: {API_README_PATH.relative_to(ROOT_DIR)}")
+336
View File
@@ -0,0 +1,336 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import itertools
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Literal
import regex as re
import yaml
logger = logging.getLogger("mkdocs")
ROOT_DIR = Path(__file__).parent.parent.parent.parent
ROOT_DIR_RELATIVE = "../../../../.."
EXAMPLE_DIR = ROOT_DIR / "examples"
EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
NAV_FILE = ROOT_DIR / "docs/.nav.yml"
def fix_case(text: str) -> str:
subs = {
"api": "API",
"cli": "CLI",
"cpu": "CPU",
"llm": "LLM",
"mae": "MAE",
"tpu": "TPU",
"gguf": "GGUF",
"lora": "LoRA",
"rlhf": "RLHF",
"vllm": "vLLM",
"openai": "OpenAI",
"lmcache": "LMCache",
"multilora": "MultiLoRA",
"mlpspeculator": "MLPSpeculator",
r"fp\d+": lambda x: x.group(0).upper(), # e.g. fp16, fp32
r"int\d+": lambda x: x.group(0).upper(), # e.g. int8, int16
}
for pattern, repl in subs.items():
text = re.sub(rf"\b{pattern}\b", repl, text, flags=re.IGNORECASE)
return text
@dataclass
class Example:
"""
Example class for generating documentation content from a given path.
Attributes:
path (Path): The path to the main directory or file.
category (str): The category of the document.
main_file (Path): The main file in the directory.
other_files (list[Path]): list of other files in the directory.
title (str): The title of the document.
Methods:
__post_init__(): Initializes the main_file, other_files, and title attributes.
determine_main_file() -> Path: Determines the main file in the given path.
determine_other_files() -> list[Path]: Determines other files in the directory excluding the main file.
determine_title() -> str: Determines the title of the document.
generate() -> str: Generates the documentation content.
""" # noqa: E501
path: Path
category: str = None
main_file: Path = field(init=False)
other_files: list[Path] = field(init=False)
title: str = field(init=False)
def __post_init__(self):
self.main_file = self.determine_main_file()
self.other_files = self.determine_other_files()
self.title = self.determine_title()
@property
def is_code(self) -> bool:
return self.main_file.suffix != ".md"
def determine_main_file(self) -> Path:
"""
Determines the main file in the given path.
If the path is a file, it returns the path itself. Otherwise, it searches
for Markdown files (*.md) in the directory and returns the first one found.
Returns:
Path: The main file path, either the original path if it's a file or the first
Markdown file found in the directory.
Raises:
IndexError: If no Markdown files are found in the directory.
""" # noqa: E501
return self.path if self.path.is_file() else list(self.path.glob("*.md")).pop()
def determine_other_files(self) -> list[Path]:
"""
Determine other files in the directory excluding the main file.
This method checks if the given path is a file. If it is, it returns an empty list.
Otherwise, it recursively searches through the directory and returns a list of all
files that are not the main file.
Returns:
list[Path]: A list of Path objects representing the other files in the directory.
""" # noqa: E501
if self.path.is_file():
return []
# Binary file extensions to exclude
binary_extensions = {
".wav",
".mp3",
".mp4",
".avi",
".mov",
".mkv", # Audio/Video
".png",
".jpg",
".jpeg",
".gif",
".bmp",
".ico",
".svg", # Images
".pdf",
".zip",
".tar",
".gz",
".bz2",
".xz", # Archives/Documents
".exe",
".so",
".dll",
".dylib", # Binaries
".bin",
".dat",
".db",
".sqlite", # Data files
}
def is_other_file(file: Path) -> bool:
return file.is_file() and file != self.main_file and file.suffix.lower() not in binary_extensions
return [file for file in self.path.rglob("*") if is_other_file(file)]
def determine_title(self) -> str:
if not self.is_code:
# Specify encoding for building on Windows
with open(self.main_file, encoding="utf-8") as f:
first_line = f.readline().strip()
match = re.match(r"^#\s+(?P<title>.+)$", first_line)
if match:
return match.group("title")
return fix_case(self.path.stem.replace("_", " ").title())
def fix_relative_links(self, content: str) -> str:
"""
Fix relative links in markdown content by converting them to gh-file
format.
Args:
content (str): The markdown content to process
Returns:
str: Content with relative links converted to gh-file format
"""
# Regex to match markdown links [text](relative_path)
# This matches links that don't start with http, https, ftp, or #
link_pattern = r"\[([^\]]*)\]\((?!(?:https?|ftp)://|#)([^)]+)\)"
def replace_link(match):
link_text = match.group(1)
relative_path = match.group(2)
# Make relative to repo root
gh_file = (self.main_file.parent / relative_path).resolve()
gh_file = gh_file.relative_to(ROOT_DIR)
# Make GitHub URL
url = "https://github.com/vllm-project/vllm-omni/"
url += "tree/main" if self.path.is_dir() else "blob/main"
gh_url = f"{url}/{gh_file}"
return f"[{link_text}]({gh_url})"
return re.sub(link_pattern, replace_link, content)
def generate(self) -> str:
content = f"# {self.title}\n\n"
url = "https://github.com/vllm-project/vllm-omni/"
url += "tree/main" if self.path.is_dir() else "blob/main"
content += f"Source <{url}/{self.path.relative_to(ROOT_DIR)}>.\n\n"
# Use long code fence to avoid issues with
# included files containing code fences too
code_fence = "``````"
if self.is_code:
main_file_rel = self.main_file.relative_to(ROOT_DIR)
content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel}"\n{code_fence}\n'
else:
with open(self.main_file, encoding="utf-8") as f:
# Skip the title from md snippets as it's been included above
main_content = f.readlines()[1:]
content += self.fix_relative_links("".join(main_content))
content += "\n"
if not self.other_files:
return content
content += "## Example materials\n\n"
for file in sorted(self.other_files):
content += f'??? abstract "{file.relative_to(self.path)}"\n'
if file.suffix != ".md":
content += f" {code_fence}{file.suffix[1:]}\n"
content += f' --8<-- "{file.relative_to(ROOT_DIR)}"\n'
if file.suffix != ".md":
content += f" {code_fence}\n"
return content
def update_nav_file(examples: list[Example]):
"""
Update the .nav.yml file to include all generated examples.
This function completely regenerates the examples section based on the actual
folder structure, ensuring consistency between the examples folder and nav file.
Args:
examples: List of Example objects that have been generated
"""
if not NAV_FILE.exists():
logger.warning("Navigation file not found: %s", NAV_FILE)
return
# Read the current nav file
with open(NAV_FILE, encoding="utf-8") as f:
nav_data = yaml.safe_load(f) or {}
nav_list = nav_data.get("nav", [])
# Find the "User Guide" section
user_guide_idx = None
examples_idx = None
for i, item in enumerate(nav_list):
if isinstance(item, dict) and "User Guide" in item:
user_guide_idx = i
user_guide_content = item["User Guide"]
# Find the "Examples" subsection
for j, subitem in enumerate(user_guide_content):
if isinstance(subitem, dict) and "Examples" in subitem:
examples_idx = j
break
break
if user_guide_idx is None or examples_idx is None:
logger.warning("Could not find 'User Guide' -> 'Examples' section in nav file")
return
# Get existing Examples section to preserve non-example items (like README.md)
existing_examples_content = nav_list[user_guide_idx]["User Guide"][examples_idx]["Examples"]
# Preserve string items (like "examples/README.md") that are not example categories
preserved_items = [
item
for item in existing_examples_content
if isinstance(item, str) and not item.startswith("user_guide/examples/")
]
# Group examples by category
examples_by_category = {}
for example in examples:
category = example.category
if category not in examples_by_category:
examples_by_category[category] = []
examples_by_category[category].append(example)
# Build the new Examples section - start with preserved items
examples_section = preserved_items.copy()
# Add examples grouped by category, sorted by category name
for category in sorted(examples_by_category.keys()):
category_examples = sorted(examples_by_category[category], key=lambda e: e.path.stem)
category_items = []
for example in category_examples:
doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
rel_path = doc_path.relative_to(ROOT_DIR / "docs")
category_items.append({example.title: str(rel_path)})
if category_items:
# Format category name (e.g., "offline_inference" -> "Offline Inference")
category_title = fix_case(category.replace("_", " ").title())
examples_section.append({category_title: category_items})
# Update the nav structure
nav_list[user_guide_idx]["User Guide"][examples_idx]["Examples"] = examples_section
# Write back to file
nav_data["nav"] = nav_list
with open(NAV_FILE, "w", encoding="utf-8") as f:
yaml.dump(nav_data, f, default_flow_style=False, sort_keys=False, allow_unicode=True)
logger.info("Updated navigation file: %s", NAV_FILE.relative_to(ROOT_DIR))
def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
logger.info("Generating example documentation")
logger.debug("Root directory: %s", ROOT_DIR.resolve())
logger.debug("Example directory: %s", EXAMPLE_DIR.resolve())
logger.debug("Example document directory: %s", EXAMPLE_DOC_DIR.resolve())
# Create the EXAMPLE_DOC_DIR if it doesn't exist
if not EXAMPLE_DOC_DIR.exists():
EXAMPLE_DOC_DIR.mkdir(parents=True)
categories = sorted(p for p in EXAMPLE_DIR.iterdir() if p.is_dir())
examples = []
glob_patterns = ["*.py", "*.md", "*.sh"]
# Find categorised examples
for category in categories:
globs = [category.glob(pattern) for pattern in glob_patterns]
for path in itertools.chain(*globs):
examples.append(Example(path, category.stem))
# Find examples in subdirectories
for path in category.glob("*/*.md"):
examples.append(Example(path.parent, category.stem))
# Generate the example documentation
for example in sorted(examples, key=lambda e: e.path.stem):
example_name = f"{example.path.stem}.md"
doc_path = EXAMPLE_DOC_DIR / example.category / example_name
if not doc_path.parent.exists():
doc_path.parent.mkdir(parents=True)
# Specify encoding for building on Windows
with open(doc_path, "w+", encoding="utf-8") as f:
f.write(example.generate())
logger.debug("Example generated: %s", doc_path.relative_to(ROOT_DIR))
# Update the navigation file
update_nav_file(examples)
+121
View File
@@ -0,0 +1,121 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This is basically a port of MyST parser’s external URL resolution mechanism
(https://myst-parser.readthedocs.io/en/latest/syntax/cross-referencing.html#customising-external-url-resolution)
to work with MkDocs.
It allows Markdown authors to use GitHub shorthand links like:
- [Text](gh-issue:123)
- <gh-pr:456>
- [File](gh-file:path/to/file.py#L10)
These are automatically rewritten into fully qualified GitHub URLs pointing to
issues, pull requests, files, directories, or projects in the
`vllm-project/vllm-omni` repository.
The goal is to simplify cross-referencing common GitHub resources
in project docs.
"""
import regex as re
from mkdocs.config.defaults import MkDocsConfig
from mkdocs.structure.files import Files
from mkdocs.structure.pages import Page
def on_page_markdown(markdown: str, *, page: Page, config: MkDocsConfig, files: Files) -> str:
"""
Custom MkDocs plugin hook to rewrite special GitHub reference links
in Markdown.
This function scans the given Markdown content for specially formatted
GitHub shorthand links, such as:
- `[Link text](gh-issue:123)`
- `<gh-pr:456>`
And rewrites them into fully-qualified GitHub URLs with GitHub icons:
- `[:octicons-mark-github-16: Link text](https://github.com/vllm-project/vllm/issues/123)`
- `[:octicons-mark-github-16: Pull Request #456](https://github.com/vllm-project/vllm/pull/456)`
Supported shorthand types:
- `gh-issue`
- `gh-pr`
- `gh-project`
- `gh-dir`
- `gh-file`
Args:
markdown (str): The raw Markdown content of the page.
page (Page): The MkDocs page object being processed.
config (MkDocsConfig): The MkDocs site configuration.
files (Files): The collection of files in the MkDocs build.
Returns:
str: The updated Markdown content with GitHub shorthand links replaced.
"""
gh_icon = ":octicons-mark-github-16:"
gh_url = "https://github.com"
repo_url = f"{gh_url}/vllm-project/vllm-omni"
org_url = f"{gh_url}/orgs/vllm-project"
# Mapping of shorthand types to their corresponding GitHub base URLs
urls = {
"issue": f"{repo_url}/issues",
"pr": f"{repo_url}/pull",
"project": f"{org_url}/projects",
"dir": f"{repo_url}/tree/main",
"file": f"{repo_url}/blob/main",
}
# Default title prefixes for auto links
titles = {
"issue": "Issue #",
"pr": "Pull Request #",
"project": "Project #",
"dir": "",
"file": "",
}
# Regular expression to match GitHub shorthand links
scheme = r"gh-(?P<type>.+?):(?P<path>.+?)(#(?P<fragment>.+?))?"
inline_link = re.compile(r"\[(?P<title>[^\[]+?)\]\(" + scheme + r"\)")
auto_link = re.compile(f"<{scheme}>")
def replace_inline_link(match: re.Match) -> str:
"""
Replaces a matched inline-style GitHub shorthand link
with a full Markdown link.
Example:
[My issue](gh-issue:123) → [:octicons-mark-github-16: My issue](https://github.com/vllm-project/vllm/issues/123)
"""
url = f"{urls[match.group('type')]}/{match.group('path')}"
if fragment := match.group("fragment"):
url += f"#{fragment}"
return f"[{gh_icon} {match.group('title')}]({url})"
def replace_auto_link(match: re.Match) -> str:
"""
Replaces a matched autolink-style GitHub shorthand
with a full Markdown link.
Example:
<gh-pr:456> → [:octicons-mark-github-16: Pull Request #456](https://github.com/vllm-project/vllm/pull/456)
"""
type = match.group("type")
path = match.group("path")
title = f"{titles[type]}{path}"
url = f"{urls[type]}/{path}"
if fragment := match.group("fragment"):
url += f"#{fragment}"
return f"[{gh_icon} {title}]({url})"
# Replace both inline and autolinks
markdown = inline_link.sub(replace_inline_link, markdown)
markdown = auto_link.sub(replace_auto_link, markdown)
return markdown
@@ -0,0 +1,47 @@
/**
* edit_and_feedback.js
*
* Enhances MkDocs Material docs pages by:
*
* 1. Adding a "Question? Give us feedback" link
* below the "Edit" button.
*
* - The link opens a GitHub issue with a template,
* auto-filled with the current page URL and path.
*
* 2. Ensuring the edit button opens in a new tab
* with target="_blank" and rel="noopener".
*/
document.addEventListener("DOMContentLoaded", function () {
const url = window.location.href;
const page = document.body.dataset.mdUrl || location.pathname;
const feedbackLink = document.createElement("a");
feedbackLink.href = `https://github.com/vllm-project/vllm-omni/issues/new?template=100-documentation.yml&title=${encodeURIComponent(
`[Docs] Feedback for \`${page}\``
)}&body=${encodeURIComponent(`📄 **Reference:**\n${url}\n\n📝 **Feedback:**\n_Your response_`)}`;
feedbackLink.target = "_blank";
feedbackLink.rel = "noopener";
feedbackLink.title = "Provide feedback";
feedbackLink.className = "md-content__button";
feedbackLink.innerHTML = `
<svg
xmlns="http://www.w3.org/2000/svg"
height="24px"
viewBox="0 -960 960 960"
width="24px"
fill="currentColor"
>
<path d="M280-280h280v-80H280v80Zm0-160h400v-80H280v80Zm0-160h400v-80H280v80Zm-80 480q-33 0-56.5-23.5T120-200v-560q0-33 23.5-56.5T200-840h560q33 0 56.5 23.5T840-760v560q0 33-23.5 56.5T760-120H200Zm0-80h560v-560H200v560Zm0-560v560-560Z"/>
</svg>
`;
const editButton = document.querySelector('.md-content__button[href*="edit"]');
if (editButton && editButton.parentNode) {
editButton.insertAdjacentElement("beforebegin", feedbackLink);
editButton.setAttribute("target", "_blank");
editButton.setAttribute("rel", "noopener");
}
});
+20
View File
@@ -0,0 +1,20 @@
// Enables MathJax rendering
window.MathJax = {
tex: {
inlineMath: [["\\(", "\\)"]],
displayMath: [["\\[", "\\]"]],
processEscapes: true,
processEnvironments: true
},
options: {
ignoreHtmlClass: ".*|",
processHtmlClass: "arithmatex"
}
};
document$.subscribe(() => {
MathJax.startup.output.clearCache()
MathJax.typesetClear()
MathJax.texReset()
MathJax.typesetPromise()
})
+21
View File
@@ -0,0 +1,21 @@
// Initialize Mermaid for diagram rendering
mermaid.initialize({
startOnLoad: false,
theme: 'default',
securityLevel: 'loose',
flowchart: {
useMaxWidth: true,
htmlLabels: true
}
});
// Render Mermaid diagrams when page content is ready
document$.subscribe(() => {
const mermaidElements = document.querySelectorAll('.mermaid');
if (mermaidElements.length > 0) {
mermaid.run({
querySelector: '.mermaid',
nodes: mermaidElements
});
}
});
+56
View File
@@ -0,0 +1,56 @@
/**
* slack_and_forum.js
*
* Adds a custom Slack and Forum button to the MkDocs Material header.
*
*/
window.addEventListener('DOMContentLoaded', () => {
const headerInner = document.querySelector('.md-header__inner');
if (headerInner) {
const slackButton = document.createElement('button');
slackButton.className = 'slack-button';
slackButton.title = 'Join us on Slack';
slackButton.style.border = 'none';
slackButton.style.background = 'transparent';
slackButton.style.cursor = 'pointer';
slackButton.innerHTML = `
<img src="https://a.slack-edge.com/80588/marketing/img/icons/icon_slack_hash_colored.png"
style="height: 1.1rem;"
alt="Slack">
`;
slackButton.addEventListener('click', () => {
window.open('https://slack.vllm.ai', '_blank', 'noopener');
});
const forumButton = document.createElement('button');
forumButton.className = 'forum-button';
forumButton.title = 'Join the Forum';
forumButton.style.border = 'none';
forumButton.style.background = 'transparent';
forumButton.style.cursor = 'pointer';
forumButton.innerHTML = `
<svg
xmlns="http://www.w3.org/2000/svg"
viewBox="0 -960 960 960"
fill="currentColor"
>
<path d="M817.85-198.15 698.46-317.54H320q-24.48 0-41.47-16.99T261.54-376v-11.69h424.61q25.39 0 43.47-18.08 18.07-18.08 18.07-43.46v-268.92h11.69q24.48 0 41.47 16.99 17 16.99 17 41.47v461.54ZM179.08-434.69l66.84-66.85h363.31q10.77 0 17.69-6.92 6.93-6.92 6.93-17.69v-246.77q0-10.77-6.93-17.7-6.92-6.92-17.69-6.92H203.69q-10.77 0-17.69 6.92-6.92 6.93-6.92 17.7v338.23Zm-36.93 89.46v-427.69q0-25.39 18.08-43.46 18.08-18.08 43.46-18.08h405.54q25.39 0 43.46 18.08 18.08 18.07 18.08 43.46v246.77q0 25.38-18.08 43.46-18.07 18.07-43.46 18.07H261.54L142.15-345.23Zm36.93-180.92V-797.54v271.39Z"/>
</svg>
`;
forumButton.addEventListener('click', () => {
window.open('https://discuss.vllm.ai/', '_blank', 'noopener');
});
const githubSource = document.querySelector('.md-header__source');
if (githubSource) {
githubSource.parentNode.insertBefore(slackButton, githubSource.nextSibling);
githubSource.parentNode.insertBefore(forumButton, slackButton.nextSibling);
}
}
});
+1
View File
@@ -0,0 +1 @@
{% extends "base.html" %}
@@ -0,0 +1,21 @@
<!-- Enables the use of toc_depth in document frontmatter https://github.com/squidfunk/mkdocs-material/issues/4827#issuecomment-1869812019 -->
<li class="md-nav__item">
<a href="{{ toc_item.url }}" class="md-nav__link">
<span class="md-ellipsis">
{{ toc_item.title }}
</span>
</a>
<!-- Table of contents list -->
{% if toc_item.children %}
<nav class="md-nav" aria-label="{{ toc_item.title | striptags }}">
<ul class="md-nav__list">
{% for toc_item in toc_item.children %}
{% if not page.meta.toc_depth or toc_item.level <= page.meta.toc_depth %}
{% include "partials/toc-item.html" %}
{% endif %}
{% endfor %}
</ul>
</nav>
{% endif %}
</li>
+221
View File
@@ -0,0 +1,221 @@
/* Warning for latest docs */
.md-banner {
background-color: var(--md-warning-bg-color);
color: var(--md-warning-fg-color);
}
/* https://christianoliff.com/blog/styling-external-links-with-an-icon-in-css/ */
a:not(:has(svg)):not(.md-icon):not(.autorefs-external) {
align-items: center;
&[href^="//"]::after,
&[href^="http://"]::after,
&[href^="https://"]::after {
content: "";
width: 12px;
height: 12px;
margin-left: 4px;
background-image: url("data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='16' height='16' stroke='gray' viewBox='0 0 16 16'%3E%3Cpath fill-rule='evenodd' d='M8.636 3.5a.5.5 0 0 0-.5-.5H1.5A1.5 1.5 0 0 0 0 4.5v10A1.5 1.5 0 0 0 1.5 16h10a1.5 1.5 0 0 0 1.5-1.5V7.864a.5.5 0 0 0-1 0V14.5a.5.5 0 0 1-.5.5h-10a.5.5 0 0 1-.5-.5v-10a.5.5 0 0 1 .5-.5h6.636a.5.5 0 0 0 .5-.5z'/%3E%3Cpath fill-rule='evenodd' d='M16 .5a.5.5 0 0 0-.5-.5h-5a.5.5 0 0 0 0 1h3.793L6.146 9.146a.5.5 0 1 0 .708.708L15 1.707V5.5a.5.5 0 0 0 1 0v-5z'/%3E%3C/svg%3E");
background-position: center;
background-repeat: no-repeat;
background-size: contain;
display: inline-block;
}
}
a[href*="localhost"]::after,
a[href*="127.0.0.1"]::after,
/* Hide external link icons for all links */
a[href^="//"]::after,
a[href^="http://"]::after,
a[href^="https://"]::after {
display: none !important;
}
/* Light mode: darker section titles */
body[data-md-color-scheme="default"] .md-nav__item--section > label.md-nav__link .md-ellipsis {
color: rgba(0, 0, 0, 0.7) !important;
font-weight: 700;
}
/* Dark mode: lighter gray section titles */
body[data-md-color-scheme="slate"] .md-nav__item--section > label.md-nav__link .md-ellipsis {
color: rgba(255, 255, 255, 0.75) !important;
font-weight: 700;
}
/* Custom admonitions */
:root {
--md-admonition-icon--announcement: url('data:image/svg+xml;charset=utf-8,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16" width="16" height="16"><path d="M3.25 9a.75.75 0 0 1 .75.75c0 2.142.456 3.828.733 4.653a.122.122 0 0 0 .05.064.212.212 0 0 0 .117.033h1.31c.085 0 .18-.042.258-.152a.45.45 0 0 0 .075-.366A16.743 16.743 0 0 1 6 9.75a.75.75 0 0 1 1.5 0c0 1.588.25 2.926.494 3.85.293 1.113-.504 2.4-1.783 2.4H4.9c-.686 0-1.35-.41-1.589-1.12A16.4 16.4 0 0 1 2.5 9.75.75.75 0 0 1 3.25 9Z"></path><path d="M0 6a4 4 0 0 1 4-4h2.75a.75.75 0 0 1 .75.75v6.5a.75.75 0 0 1-.75.75H4a4 4 0 0 1-4-4Zm4-2.5a2.5 2.5 0 1 0 0 5h2v-5Z"></path><path d="M15.59.082A.75.75 0 0 1 16 .75v10.5a.75.75 0 0 1-1.189.608l-.002-.001h.001l-.014-.01a5.775 5.775 0 0 0-.422-.25 10.63 10.63 0 0 0-1.469-.64C11.576 10.484 9.536 10 6.75 10a.75.75 0 0 1 0-1.5c2.964 0 5.174.516 6.658 1.043.423.151.787.302 1.092.443V2.014c-.305.14-.669.292-1.092.443C11.924 2.984 9.713 3.5 6.75 3.5a.75.75 0 0 1 0-1.5c2.786 0 4.826-.484 6.155-.957.665-.236 1.154-.47 1.47-.64.144-.077.284-.161.421-.25l.014-.01a.75.75 0 0 1 .78-.061Z"></path></svg>');
--md-admonition-icon--important: url('data:image/svg+xml;charset=utf-8,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16" width="16" height="16"><path d="M4.47.22A.749.749 0 0 1 5 0h6c.199 0 .389.079.53.22l4.25 4.25c.141.14.22.331.22.53v6a.749.749 0 0 1-.22.53l-4.25 4.25A.749.749 0 0 1 11 16H5a.749.749 0 0 1-.53-.22L.22 11.53A.749.749 0 0 1 0 11V5c0-.199.079-.389.22-.53Zm.84 1.28L1.5 5.31v5.38l3.81 3.81h5.38l3.81-3.81V5.31L10.69 1.5ZM8 4a.75.75 0 0 1 .75.75v3.5a.75.75 0 0 1-1.5 0v-3.5A.75.75 0 0 1 8 4Zm0 8a1 1 0 1 1 0-2 1 1 0 0 1 0 2Z"></path></svg>');
--md-admonition-icon--code: url('data:image/svg+xml;charset=utf-8,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16"><path d="m11.28 3.22 4.25 4.25a.75.75 0 0 1 0 1.06l-4.25 4.25a.749.749 0 0 1-1.275-.326.75.75 0 0 1 .215-.734L13.94 8l-3.72-3.72a.749.749 0 0 1 .326-1.275.75.75 0 0 1 .734.215m-6.56 0a.75.75 0 0 1 1.042.018.75.75 0 0 1 .018 1.042L2.06 8l3.72 3.72a.749.749 0 0 1-.326 1.275.75.75 0 0 1-.734-.215L.47 8.53a.75.75 0 0 1 0-1.06Z"/></svg>');
--md-admonition-icon--console: url('data:image/svg+xml;charset=utf-8,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16"><path d="M0 2.75C0 1.784.784 1 1.75 1h12.5c.966 0 1.75.784 1.75 1.75v10.5A1.75 1.75 0 0 1 14.25 15H1.75A1.75 1.75 0 0 1 0 13.25Zm1.75-.25a.25.25 0 0 0-.25.25v10.5c0 .138.112.25.25.25h12.5a.25.25 0 0 0 .25-.25V2.75a.25.25 0 0 0-.25-.25ZM7.25 8a.75.75 0 0 1-.22.53l-2.25 2.25a.749.749 0 0 1-1.275-.326.75.75 0 0 1 .215-.734L5.44 8 3.72 6.28a.749.749 0 0 1 .326-1.275.75.75 0 0 1 .734.215l2.25 2.25c.141.14.22.331.22.53m1.5 1.5h3a.75.75 0 0 1 0 1.5h-3a.75.75 0 0 1 0-1.5"/></svg>');
}
.md-typeset .admonition.announcement,
.md-typeset details.announcement {
border-color: rgb(255, 110, 66);
}
.md-typeset .admonition.important,
.md-typeset details.important {
border-color: rgb(239, 85, 82);
}
.md-typeset .admonition.code,
.md-typeset details.code {
border-color: #64dd17
}
.md-typeset .admonition.console,
.md-typeset details.console {
border-color: #64dd17
}
.md-typeset .announcement > .admonition-title,
.md-typeset .announcement > summary {
background-color: rgb(255, 110, 66, 0.1);
}
.md-typeset .important > .admonition-title,
.md-typeset .important > summary {
background-color: rgb(239, 85, 82, 0.1);
}
.md-typeset .code > .admonition-title,
.md-typeset .code > summary {
background-color: #64dd171a;
}
.md-typeset .console > .admonition-title,
.md-typeset .console > summary {
background-color: #64dd171a;
}
.md-typeset .announcement > .admonition-title::before,
.md-typeset .announcement > summary::before {
background-color: rgb(239, 85, 82);
-webkit-mask-image: var(--md-admonition-icon--announcement);
mask-image: var(--md-admonition-icon--announcement);
}
.md-typeset .important > .admonition-title::before,
.md-typeset .important > summary::before {
background-color: rgb(239, 85, 82);
-webkit-mask-image: var(--md-admonition-icon--important);
mask-image: var(--md-admonition-icon--important);
}
.md-typeset .code > .admonition-title::before,
.md-typeset .code > summary::before {
background-color: #64dd17;
-webkit-mask-image: var(--md-admonition-icon--code);
mask-image: var(--md-admonition-icon--code);
}
.md-typeset .console > .admonition-title::before,
.md-typeset .console > summary::before {
background-color: #64dd17;
-webkit-mask-image: var(--md-admonition-icon--console);
mask-image: var(--md-admonition-icon--console);
}
/* Make label fully visible on hover */
.md-content__button[href*="edit"]:hover::after {
opacity: 1;
}
/* Hide edit button on generated docs/examples pages */
@media (min-width: 960px) {
.md-content__button[href*="docs/examples/"] {
display: none !important;
}
}
.md-content__button-wrapper {
position: absolute;
top: 0.6rem;
right: 0.8rem;
display: flex;
flex-direction: row;
align-items: center;
gap: 0.4rem;
z-index: 1;
}
.md-content__button-wrapper a {
display: inline-flex;
align-items: center;
justify-content: center;
height: 24px;
width: 24px;
color: var(--md-default-fg-color);
text-decoration: none;
}
.md-content__button-wrapper a:hover {
color: var(--md-accent-fg-color);
}
/* Slack and Forum css */
.slack-button,
.forum-button {
display: inline-flex;
align-items: center;
justify-content: center;
margin-left: 0.4rem;
height: 24px;
}
.slack-button img {
height: 18px;
filter: none !important;
}
.slack-button:hover,
.forum-button:hover {
opacity: 0.7;
}
.forum-button svg {
height: 28px;
opacity: 0.9;
transform: translateY(2px);
}
/* For logo css */
[data-md-color-scheme="default"] .logo-dark {
display: none;
}
[data-md-color-scheme="slate"] .logo-light {
display: none;
}
/* Outline for content tabs */
.md-typeset .tabbed-set {
border: 0.075rem solid var(--md-default-fg-color);
border-radius: 0.2rem;
}
.md-typeset .tabbed-content {
padding: 0 0.6em;
}
/* Hide link icon in header logo */
.md-header__button.md-logo :is(img, svg) {
pointer-events: none;
}
.md-header__button.md-logo::after {
display: none !important;
}
/* Hide link icons in content tabs (tabbed content) */
.md-typeset .tabbed-set > label::after,
.md-typeset .tabbed-set > input:checked + label::after {
display: none !important;
}
/* Hide link icons in navigation tabs */
.md-nav__link[href]::after,
.md-nav__item--nested > .md-nav__link::after {
display: none !important;
}
/* Hide link icons in top navigation tabs */
.md-tabs__link::after {
display: none !important;
}
+61
View File
@@ -0,0 +1,61 @@
# Supported Models
vLLM-Omni supports unified multimodal comprehension and generation models across various tasks.
## Model Implementation
If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm_omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>.
## List of Supported Models for Nvidia GPU / AMD GPU
<style>
th {
white-space: nowrap;
min-width: 0 !important;
}
</style>
| Architecture | Models | Example HF Models |
|--------------|--------|-------------------|
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` |
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` |
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` |
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
|`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
| `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `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` |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` |
|`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
|`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
|`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
|`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
|`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` |
|`StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` |
|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |
## List of Supported Models for NPU
<style>
th {
white-space: nowrap;
min-width: 0 !important;
}
</style>
| Architecture | Models | Example HF Models |
|--------------|--------|-------------------|
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` |
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` |
|`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
+249
View File
@@ -0,0 +1,249 @@
# Image Generation API
vLLM-Omni provides an OpenAI DALL-E compatible API for text-to-image generation using diffusion models.
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
## Quick Start
### Start the Server
For example...
```bash
# Qwen-Image
vllm serve Qwen/Qwen-Image --omni --port 8000
# Z-Image Turbo
vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8000
```
### Generate Images
**Using curl:**
```bash
curl -X POST http://localhost:8000/v1/images/generations \
-H "Content-Type: application/json" \
-d '{
"prompt": "a dragon laying over the spine of the Green Mountains of Vermont",
"size": "1024x1024",
"seed": 42
}' | jq -r '.data[0].b64_json' | base64 -d > dragon.png
```
**Using Python:**
```python
import requests
import base64
from PIL import Image
import io
response = requests.post(
"http://localhost:8000/v1/images/generations",
json={
"prompt": "a black and white cat wearing a princess tiara",
"size": "1024x1024",
"num_inference_steps": 50,
"seed": 42,
}
)
# Decode and save
img_data = response.json()["data"][0]["b64_json"]
img_bytes = base64.b64decode(img_data)
img = Image.open(io.BytesIO(img_bytes))
img.save("cat.png")
```
**Using OpenAI SDK:**
```python
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8000/v1", api_key="none")
response = client.images.generate(
model="Qwen/Qwen-Image",
prompt="a horse jumping over a fence nearby a babbling brook",
n=1,
size="1024x1024",
response_format="b64_json"
)
# Note: Extension parameters (seed, steps, cfg) require direct HTTP requests
```
## API Reference
### Endpoint
```
POST /v1/images/generations
Content-Type: application/json
```
### Request Parameters
#### OpenAI Standard Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `prompt` | string | **required** | Text description of the desired image |
| `model` | string | server's model | Model to use (optional, should match server if specified) |
| `n` | integer | 1 | Number of images to generate (1-10) |
| `size` | string | model defaults | Image dimensions in WxH format (e.g., "1024x1024", "512x512") |
| `response_format` | string | "b64_json" | Response format (only "b64_json" supported) |
| `user` | string | null | User identifier for tracking |
#### vllm-omni Extension Parameters
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `negative_prompt` | string | null | Text describing what to avoid in the image |
| `num_inference_steps` | integer | model defaults | Number of diffusion steps |
| `guidance_scale` | float | model defaults | Classifier-free guidance scale (typically 0.0-20.0) |
| `true_cfg_scale` | float | model defaults | True CFG scale (model-specific parameter, may be ignored if not supported) |
| `seed` | integer | null | Random seed for reproducibility |
### Response Format
```json
{
"created": 1701234567,
"data": [
{
"b64_json": "<base64-encoded PNG>",
"url": null,
"revised_prompt": null
}
]
}
```
## Examples
### Multiple Images
```bash
curl -X POST http://localhost:8000/v1/images/generations \
-H "Content-Type: application/json" \
-d '{
"prompt": "a steampunk city set in a valley of the Adirondack mountains",
"n": 4,
"size": "1024x1024",
"seed": 123
}'
```
This generates 4 images in a single request.
### With Negative Prompt
```python
response = requests.post(
"http://localhost:8000/v1/images/generations",
json={
"prompt": "a portrait of a skier in deep powder snow",
"negative_prompt": "blurry, low quality, distorted, ugly",
"num_inference_steps": 100,
"size": "1024x1024",
}
)
```
## Parameter Handling
The API passes parameters directly to the diffusion pipeline without model-specific transformation:
- **Default values**: When parameters are not specified, the underlying model uses its own defaults
- **Pass-through design**: User-provided values are forwarded directly to the diffusion engine
- **Minimal validation**: Only basic type checking and range validation at the API level
### Parameter Compatibility
The API passes parameters directly to the diffusion pipeline without model-specific validation.
- Unsupported parameters may be silently ignored by the model
- Incompatible values will result in errors from the underlying pipeline
- Recommended values vary by model - consult model documentation
**Best Practice:** Start with the model's recommended parameters, then adjust based on your needs.
## Error Responses
### 400 Bad Request
Invalid parameters (e.g., model mismatch):
```json
{
"detail": "Invalid size format: '1024x'. Expected format: 'WIDTHxHEIGHT' (e.g., '1024x1024')."
}
```
### 422 Unprocessable Entity
Validation errors (missing required fields):
```json
{
"detail": [
{
"loc": ["body", "prompt"],
"msg": "field required",
"type": "value_error.missing"
}
]
}
```
### 503 Service Unavailable
Diffusion engine not initialized:
```json
{
"detail": "Diffusion engine not initialized. Start server with a diffusion model."
}
```
## Troubleshooting
### Server Not Running
```bash
# Check if server is responding
curl http://localhost:8000/v1/images/generations \
-H "Content-Type: application/json" \
-d '{"prompt": "test"}'
```
### Out of Memory
If you encounter OOM errors:
1. Reduce image size: `"size": "512x512"`
2. Reduce inference steps: `"num_inference_steps": 25`
3. Generate fewer images: `"n": 1`
## Testing
Run the test suite to verify functionality:
```bash
# All image generation tests
pytest tests/entrypoints/openai_api/test_image_server.py -v
# Specific test
pytest tests/entrypoints/openai_api/test_image_server.py::test_generate_single_image -v
```
## Development
Enable debug logging to see prompts and generation details:
```bash
vllm serve Qwen/Qwen-Image --omni \
--uvicorn-log-level debug
```
Binary file not shown.

After

Width:  |  Height:  |  Size: 107 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 97 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 126 KiB

Loaded 100 of 522 files, more files were not shown because too many files have changed in this diff. Show more